#!/usr/bin/env python3

"""Create an ECR image scan job and return the results."""

import copy
import json
import re
import subprocess
import sys
import tarfile
import tempfile
from datetime import datetime
from urllib.parse import quote

import boto3
from owslogger import logger
from prettytable import PrettyTable

import config

# Global logger
log = logger.setup(
    config.LOGGER_DSN,
    config.ENVIRONMENT,
    config.LOGGER_NAME,
    config.LOGGER_LEVEL,
    config.SERVICE_NAME,
    config.SERVICE_VERSION,
)

ecr_client = boto3.client("ecr", region_name=config.AWS_REGION)
inspector_scan_client = boto3.client("inspector-scan", region_name=config.AWS_REGION)


def main():
    """Execute main entrypoint."""
    image_path = config.IMAGE_PATH

    log.info(f"Generating SBOM for image {image_path}")
    sbom_file = "sbom.json"
    inspector_sbomgen_command = [
        "./inspector-sbomgen",
        "container",
        "--image",
        image_path,
        "-o",
        sbom_file,
        "--disable-progress-bar",
        "--skip-scanners",
        config.INSPECTOR_SBOMGEN_SKIP_SCANNERS,
    ]

    try:
        subprocess.run(
            inspector_sbomgen_command,
            check=True,
        )
    except subprocess.CalledProcessError as error:
        log.error("Failed to generate SBOM. Check inspector-sbomgen output.")
        raise error

    with open(sbom_file) as sbom:
        sbom_json = json.load(sbom)

    findings = scan_sbom(sbom_json)
    log.info(f"Found {len(findings)} total findings")

    block_findings, warning_findings = get_failed_findings(findings, sbom_json)

    exit_code = 0
    if block_findings or warning_findings:
        print_findings(block_findings, warning_findings)

    if block_findings or warning_findings:
        if config.ECR_REPOSITORY_NAME and config.IMAGE_TAG:
            print_inspector_link()
        exit_code = 1 if block_findings else 2
    else:
        log.info("No failing vulnerabilities found!")
        exit_code = 0

    sys.exit(exit_code)


def calculate_days_difference(date_str):
    """
    Calculate the number of days between the current date and the given date string.

    Parameters:
        date_str (str): The date in ISO format (YYYY-MM-DDTHH:MM:SSZ).

    Returns:
        int: Number of days between the current date and the given date.
    """
    current_date = datetime.now()
    created_at = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%SZ")
    days_difference = (current_date - created_at).days
    return days_difference


def is_non_blocking_vulnerability_stale(vulnerability_findings, threshold_days):
    days_left = calculate_days_difference(vulnerability_findings["published"])

    # If the non-blocking vulnerability is older than the threshold, return True
    return days_left > threshold_days


def scan_sbom(sbom):
    """Scan an SBOM for vulnerabilities.

    The SBOM is split into batches if the number of components exceeds the limit
    of the ScanSbom Inspector API.
    """
    sbom_copy = copy.deepcopy(sbom)
    sbom_batch = create_sbom_batch(sbom_copy)
    findings = []

    while len(sbom_copy["components"]) > 0:
        if len(sbom_batch["components"]) == config.MAXIMUM_NUMBER_OF_SBOM_COMPONENTS:
            log.info(
                f"Number of components in SBOM exceeded maximum: "
                f"{config.MAXIMUM_NUMBER_OF_SBOM_COMPONENTS}. "
                f"Creating new batch."
            )
            merge_findings(findings, close_sbom_batch(sbom_batch))
            sbom_batch = create_sbom_batch(sbom_copy)
        sbom_batch["components"].append(sbom_copy["components"].pop())

    merge_findings(findings, close_sbom_batch(sbom_batch))
    return findings


def create_sbom_batch(sbom):
    sbom_batch = copy.deepcopy(sbom)
    sbom_batch["components"] = []
    return sbom_batch


def close_sbom_batch(sbom_batch):
    scan_response = inspector_scan_client.scan_sbom(
        sbom=json.dumps(sbom_batch), outputFormat="INSPECTOR"
    )
    return scan_response["sbom"].get("vulnerabilities", [])


def merge_findings(findings, findings_to_merge):
    """Merges the findings from findings_to_merge into findings."""
    for finding in findings_to_merge:
        matched_finding = next(
            ((f for f in findings if f["id"] == finding["id"])), None
        )
        if matched_finding:
            # Merge the findings
            if "affects" not in matched_finding:
                matched_finding["affects"] = []
            matched_finding["affects"].extend(finding.get("affects", []))
        else:
            # Add the new finding to the list
            findings.append(finding)


def get_failed_findings(findings, sbom):
    """Iterate over findings to return failed and warning findings.

    Findings can be suppressed for the following reasons:
    * Finding severity is not in VULNERABILITY_SEVERITIES_TO_FAIL
    * Finding is explicitly ignored by VULNERABILITIES_TO_IGNORE
    * Finding is for a package whose path matches FILE_PATHS_TO_IGNORE
    * Finding is suppressed based on an operating-system level exception defined in EXCEPTIONS_BY_OS
    * Finding is suppressed based on an exception for the image's parent image, defined in EXCEPTIONS_BY_PARENT_IMAGE
    * Finding is for the Linux kernel and there is no fix available
    * Block findings that exceed the allowed age (`DAYS_FOR_ERROR`) for their severity level and
      have an available fix or not.
    * Warning findings that do not meet the criteria for being a block finding but are still considered
      vulnerabilities.
    """
    os_info = get_os_info(sbom)

    os_name = os_info.get("name")
    os_version = os_info.get("version")

    log.info(f"OS Name: {os_name}")
    log.info(f"OS Version: {os_version}")

    os_exceptions_by_version = config.EXCEPTIONS_BY_OS.get(os_name, {})
    os_exceptions = set(os_exceptions_by_version.get(os_version, []))

    parent_images = get_parent_images()
    log.info(f"Parent images: {parent_images}")
    parent_image_exceptions = get_parent_image_exceptions(parent_images)

    block_findings = []
    warning_findings = []
    ignored_vulnerabilities = set(config.VULNERABILITIES_TO_IGNORE)
    for item in findings:
        severity = item["severity"].upper()
        vulnerability_id = item["id"]

        if severity not in config.VULNERABILITY_SEVERITIES_TO_FAIL:
            log.info(
                f"Ignoring vulnerability {vulnerability_id} because severity "
                f"{severity} is not in "
                f"{config.VULNERABILITY_SEVERITIES_TO_FAIL}"
            )
            continue

        if vulnerability_id in config.VULNERABILITIES_TO_IGNORE:
            log.info(
                f"Ignoring explicitly ignored vulnerability " f"{vulnerability_id}"
            )
            ignored_vulnerabilities.discard(vulnerability_id)
            continue

        file_paths = [package.get("path", "") for package in item.get("affects", [])]

        if file_paths and all(ignore_file_path(file_path) for file_path in file_paths):
            log.info(
                f"Ignoring vulnerability {vulnerability_id} because "
                f"vulnerable packages match list of ignored file paths. "
                f"File paths: {file_paths}"
            )
            continue

        if vulnerability_id in os_exceptions:
            log.info(
                f"Ignoring vulnerability {vulnerability_id} because there "
                f"is a global exception for the image's operating system "
                f"({os_name} {os_version})"
            )
            os_exceptions.discard(vulnerability_id)
            continue

        if vulnerability_id in parent_image_exceptions:
            log.info(
                f"Ignoring vulnerability {vulnerability_id} because there "
                f"is a global exception for the following parent image(s): "
                f"({parent_image_exceptions[vulnerability_id]})"
            )
            del parent_image_exceptions[vulnerability_id]
            continue

        if is_unfixed_kernel_vulnerability(item):
            log.info(
                f"Ignoring {vulnerability_id} as it is an unfixed kernel vulnerability"
            )
            continue

        fixed_versions = get_fixed_versions(item)
        if not fixed_versions:
            if is_non_blocking_vulnerability_stale(
                item, config.DAYS_FOR_ERROR["NON_BLOCKING"]
            ):
                block_findings.append(item)
            else:
                days_elapsed = calculate_days_difference(item["published"])
                days_left = config.DAYS_FOR_ERROR["NON_BLOCKING"] - days_elapsed
                item["days_left"] = days_left
                warning_findings.append(item)
        elif is_block_finding(item, config.DAYS_FOR_ERROR):
            block_findings.append(item)
        else:
            days_elapsed = calculate_days_difference(item["published"])
            days_left = config.DAYS_FOR_ERROR[severity] - days_elapsed
            item["days_left"] = days_left
            warning_findings.append(item)

    for ignored_vulnerability in ignored_vulnerabilities:
        log.warning(
            f"Warning: Ignored vulnerability {ignored_vulnerability} is no longer present. "
            f"Consider removing it from the ignore list."
        )

    for os_exception in os_exceptions:
        log.warning(
            f"Warning: Vulnerability {os_exception} is not present but is defined as an OS-level exception "
            f"for OS ({os_name} {os_version}). "
            f"Consider removing it from the ignore list."
        )

    for exception, parent_images in parent_image_exceptions.items():
        log.warning(
            f"Warning: Vulnerability {exception} is not present but is defined as an exception for the "
            f"following parent images: {parent_images}. "
            f"Consider removing it from the ignore list."
        )

    return block_findings, warning_findings


def get_parent_image_exceptions(parent_images):
    """Return a dict of exceptions for the given parent images.

    The key of the resulting dict is the vulnerability ID and the value is a list of parent images
    to which that applies.
    """
    parent_image_exceptions = {}

    for parent_image in parent_images:
        if parent_image not in config.EXCEPTIONS_BY_PARENT_IMAGE:
            continue

        for vulnerability_id in config.EXCEPTIONS_BY_PARENT_IMAGE[parent_image]:
            if vulnerability_id not in parent_image_exceptions:
                parent_image_exceptions[vulnerability_id] = []
            parent_image_exceptions[vulnerability_id].append(parent_image)

    return parent_image_exceptions


def is_unfixed_kernel_vulnerability(item):
    """Determine if this vulnerability is a Linux kernel vulnerability without an available fix"""
    return any(
        is_kernel_vulnerability(component) and not component.get("fixed_version")
        for component in item.get("affects", [])
    )


def is_kernel_vulnerability(component):
    """Determine if the given component is the Linux kernel"""
    installed_version = component.get("installed_version", "")
    return re.match(r"^pkg:(deb|rpm|apk)/[^/]+/linux@", installed_version)


def ignore_file_path(file_path):
    """Ignore file path helper method."""
    return any(re.match(pattern, file_path) for pattern in config.FILE_PATHS_TO_IGNORE)


def calculate_age_in_days(created_date):
    """Calculate the age in days from the creation date to today."""
    issue_date = datetime.strptime(created_date, "%Y-%m-%dT%H:%M:%SZ")
    current_date = datetime.today()
    delta = current_date - issue_date
    return delta.days


def get_installed_versions(finding):
    """Extract installed versions from a finding."""
    affects = finding.get("affects", [])
    return [
        affect.get("installed_version")
        for affect in affects
        if affect.get("installed_version")
    ]


def get_fixed_versions(finding):
    """Extract fixed versions from a finding."""
    affects = finding.get("affects", [])
    return [
        affect.get("fixed_version") for affect in affects if affect.get("fixed_version")
    ]


def get_paths(finding):
    """Extract fixed versions from a finding."""
    affects = finding.get("affects", [])
    return [affect.get("path") for affect in affects if affect.get("path")]


def is_block_finding(finding, days_for_error):
    """Determine if a finding should block the build based on its age, severity, and fixed versions."""
    severity = finding["severity"].upper()
    age_in_days = calculate_age_in_days(finding["published"])
    return age_in_days > days_for_error[severity]


def log_findings_count(findings_type, count):
    """Log the count of findings with specific formatting."""
    if findings_type == "BLOCK":
        log.info(f"{config.RED}###### BLOCK FINDINGS: {count} ######{config.RESET}")
    elif findings_type == "WARNING":
        log.info(
            f"{config.YELLOW}###### WARNING FINDINGS: {count} ######{config.RESET}"
        )


def print_findings(block_findings, warning_findings):
    """Print failed and warning findings."""

    # Define severity order (most severe first)
    severity_order = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "UNTRIAGED": 4}

    block_findings.sort(
        key=lambda x: (
            severity_order.get(x["severity"].upper(), 99),  # Most severe first
            x["id"],  # Alphabetical
        )
    )

    warning_findings.sort(
        key=lambda x: (
            severity_order.get(x["severity"].upper(), 99),  # Most severe first
            -int(x["days_left"]),  # Descending order (largest first)
            x["id"],  # Alphabetical
        )
    )

    # Create and populate the table
    findings_table = PrettyTable()
    findings_table.field_names = [
        "Vulnerability ID",
        "Installed Version",
        "Fixed Version",
        "Path",
        "Severity",
        "Blocking",
        "Grace Period",
    ]

    for finding in block_findings:
        findings_table.add_row(
            [
                finding["id"],
                next(iter(get_installed_versions(finding)), "N/A"),
                next(iter(get_fixed_versions(finding)), "N/A"),
                next(iter(get_paths(finding)), "N/A"),
                finding["severity"].upper(),
                "Yes",
                "N/A",
            ]
        )

    for finding in warning_findings:
        findings_table.add_row(
            [
                finding["id"],
                next(iter(get_installed_versions(finding)), "N/A"),
                next(iter(get_fixed_versions(finding)), "N/A"),
                next(iter(get_paths(finding)), "N/A"),
                finding["severity"].upper(),
                "No",
                finding["days_left"],
            ]
        )

    # Print finding counts
    if block_findings:
        log_findings_count("BLOCK", len(block_findings))
    if warning_findings:
        log_findings_count("WARNING", len(warning_findings))

    # Print table of findings
    log.info(findings_table)

    if config.FINDINGS_OUTPUT_FILE:
        write_findings_to_file(findings_table)


def write_findings_to_file(findings_table):
    """Write the findings table to the configured output file."""
    log.info(f"Writing findings table to {config.FINDINGS_OUTPUT_FILE}")
    with open(config.FINDINGS_OUTPUT_FILE, "w") as output_file:
        output_file.write(findings_table.get_string())


def print_inspector_link():
    """Print link to Inspector console."""
    response = ecr_client.describe_images(
        repositoryName=config.ECR_REPOSITORY_NAME,
        imageIds=[
            {
                "imageTag": config.IMAGE_TAG,
            }
        ],
    )
    image_details = response["imageDetails"][0]

    image_digest = image_details["imageDigest"]
    account_id = image_details["registryId"]
    image_arn = get_image_arn(account_id, image_digest)
    image_arn = quote(image_arn, safe="")

    aws_console_link = f"https://{config.AWS_REGION}.console.aws.amazon.com/inspector/v2/home?region={config.AWS_REGION}#/findings/container-image/{image_arn}"  # noqa: E501
    log.info(
        f"\nFor more information, see the Inspector console at " f"{aws_console_link}."
    )
    if account_id != config.PROD_ACCOUNT_ID:
        assume_role_link = f"https://signin.aws.amazon.com/switchrole?account={account_id}&roleName=generic-engineer-role"  # noqa: E501
        log.info(
            f"\nTo view the information in the Inspector console, you will "
            f"need to assume the appropriate role in account {account_id}. "
            f"You can do this using this link: {assume_role_link}."
        )


def get_image_arn(account_id, digest):
    """Get image arn."""
    return f"arn:aws:ecr:{config.AWS_REGION}:{account_id}:repository/{config.ECR_REPOSITORY_NAME}/{digest}"


def get_os_info(sbom_json):
    """Return the operating system info for the given SBOM."""
    components = sbom_json.get("components", [])
    os_components = (c for c in components if c["type"] == "operating-system")
    os_component = next(os_components, {})
    return os_component


def get_parent_images():
    """Determine the image's parent image(s)."""
    labels = get_image_labels()
    return labels.get("parent-images", "").split(",")


def get_image_labels():
    """Return the Docker labels for the image."""
    image_config = get_image_config()
    return image_config.get("Labels", {}) or {}


def get_image_config():
    """Return the Docker image configuration from the tarball."""
    with tempfile.TemporaryDirectory() as tmp:
        with tarfile.open(config.IMAGE_PATH) as tar:
            tar.extractall(path=tmp)
        with open(f"{tmp}/manifest.json") as manifest:
            manifest_json = json.load(manifest)
            image_config_file = manifest_json[0]["Config"]
        with open(f"{tmp}/{image_config_file}") as image_config:
            image_config_json = json.load(image_config)
            return image_config_json["config"]


if __name__ == "__main__":
    main()