"""Script to parse results from Trivy scan.

1) Reads in results from two files (json and table format) generated by Trivy.
2) If a GitHub issue for the current release does not exist yet, creates a
GitHub issue in the given repo if any vulnerabilities were found.
3) If a GitHub issue for the current release exists already, either overwrite
this issue with the new results, or close this issue if the new scan returned
no vulnerabilities.
3) Fails the CI build if any vulerabilities of HIGH/CRITICAL severity exist.

Both the json and table Trivy results must be created prior to running
this script. These can be generated with:
    trivy image --format json --light --no-progress -o <trivy_output>.json
        gcr.io/cre-tools/ts-bridge
    trivy image --light  --no-progress -o <trivy_output>.table
        gcr.io/cre-tools/ts-bridge
"""
import json
import re
import sys
import textwrap
from absl import app
from absl import flags
from github import BadAttributeException
from github import Github
from github import GithubException
from pathlib import Path

FLAGS = flags.FLAGS
flags.DEFINE_string(
    "build_id", None,
    "ID of the current Cloud Build, specified using $BUILD_ID substitution")
flags.DEFINE_string(
    "commit_id", None,
    "ID of the current commit, specified using $COMMIT_ID substitution")
flags.DEFINE_string(
    "release_tag", None, "Number of the latest release/tag. "
    "This can be retrieved using \n "
    "git describe --abbrev=0 --tags > _release_tag")
flags.DEFINE_string(
    "repo_name", None,
    "Name of the github repo where potential issues will be created.")
flags.DEFINE_string(
    "token_file", None,
    "Name of the file where the GitHub token has been stored.\n"
    "This can be created using \n"
    "gcloud secrets versions access latest "
    "--secret=Ts-bridge-bot-token --format='get(payload.data)' | "
    "tr '_-' '/+' | base64 -d > git_token.txt")
flags.DEFINE_string("trivy_file", None,
                    "Name of the files where results from Trivy are stored")


def validate_files():
    """Validate that the files given through flags exists in the curr dir."""
    error = None
    if not Path("{}.json".format(FLAGS.trivy_file)).is_file():
        error = ("Please run \n"
                 "trivy image --format json --light --no-progress -o {}.json "
                 "gcr.io/cre-tools/ts-bridge \n").format(FLAGS.trivy_file)
    elif not Path("{}.table".format(FLAGS.trivy_file)).is_file():
        error = ("Please run \n"
                 "trivy image --light --no-progress -o {}.table "
                 "gcr.io/cre-tools/ts-bridge \n").format(FLAGS.trivy_file)
    elif not Path(FLAGS.token_file).is_file():
        error = ("Please run \n"
                 "gcloud secrets versions access latest "
                 "--secret=Ts-bridge-bot-token --format='get(payload.data)' | "
                 "tr '_-' '/+' | base64 -d > {}").format(FLAGS.token_file)
    if error:
        sys.exit(error)


def load_results():
    """Load the results from Trivy."""
    trivy_out_json = "{}.json".format(FLAGS.trivy_file)
    trivy_out_table = "{}.table".format(FLAGS.trivy_file)
    with open(trivy_out_json) as f:
        # trivy_json is a dictionary of the JSON result used for parsing
        # vulnerabilities.
        # Index is 0 because there is only one target built.
        trivy_json = json.load(f)[0]
    with open(trivy_out_table) as f:
        # trivy_table is the results in a more readable table format which can
        # be included in GitHub issues.
        trivy_table = f.read()
    return [trivy_json, trivy_table]


def get_severity_list(vulnerabilities):
    """Filters out all the severities in a given list of vulnerabilities."""
    severity_list = [v.get("Severity") for v in vulnerabilities]
    return sorted(set(severity_list))


def high_or_critical_exists(severity_list):
    return "HIGH" in severity_list or "CRITICAL" in severity_list


def get_github_repo():
    """Connects to GitHub API and returns a repo object."""
    with open(FLAGS.token_file) as f:
        token = f.read()
    github = Github(token)
    try:
        repo = github.get_repo(FLAGS.repo_name)
    except (BadAttributeException, GithubException) as e:
        sys.exit(
            ("Failed to get repo with name {} due to an exception from GitHub."
             "\nThe error returned by GitHub API was {}").format(
                FLAGS.repo_name, e))

    return repo


def build_issue(num_vulnerabilities, severity_list, table, json_output):
    """Creates the title and body of the GitHub issue."""
    title = ("Vulnerability found in release {}").format(FLAGS.release_tag)

    body = ("Trivy has detected {} vulnerabilities in your latest "
            "build.").format(num_vulnerabilities)

    if high_or_critical_exists(severity_list):
        title += ": Images from commit {} cannot be released".format(
            FLAGS.commit_id)

        body += " Please correct this issue so the new images can be "
        "published on Container Registry."

    body_args = dict(build_id=FLAGS.build_id, commit_id=FLAGS.commit_id,
                     release_tag=FLAGS.release_tag, target_name=json_output["Target"],
                     table=table, json_output=json.dumps(json_output, sort_keys=True, indent=4))
    body += textwrap.dedent("""
        \n
        **Cloud Build ID:** {build_id}
        **Commit ID:** {commit_id}
        **Release Tag:** {release_tag}
        **Target:** {target_name}
        ```
        {table}
        ```

        <details>
          <summary>JSON output from Trivy</summary>

        ```json
        {json_output}
        ```
        </details>""").format(**body_args)

    return (title, body)


def create_or_update_issue(repo, existing_issue, title, body):
    """Creates or updates a Github issue for the current release."""

    try:
        if existing_issue:
            existing_issue.edit(title=title, body=body)
            return existing_issue.number

        new_issue = repo.create_issue(title=title, body=body)
        return new_issue.number
    except (BadAttributeException, GithubException) as e:
        sys.exit(
            ("Failed to create or update issue due to an exception from GitHub."
             "\nThe error returned by GitHub API was {}").format(e))


def get_duplicate_issue(repo):
    """Finds a GitHub issue for the current release, else return None"""
    issues = repo.get_issues(creator="ts-bridge-bot",
                             state="open", sort="updated")

    # Filter out the issue with the same release number as current build
    title = ("Vulnerability found in release {}").format(FLAGS.release_tag)
    prev_issue = [issue for issue in issues if issue.title.startswith(title)]

    if prev_issue:
        if len(prev_issue) > 1:
            print("WARNING: More than one vulnerability issue exists for this "
                  "release. We will work with the most recently updated issue.")
        return prev_issue[0]
    else:
        return None


def process_vulnerabilities(repo, existing_issue, trivy_json, trivy_table):
    """Updates issue and build log based on non-empty vulnerabilities.

    Converts the Trivy output into a formatted title and body for a GitHub
    Issue, then creates or updates this issue on GitHub. Also produces a build
    log message and fails the build if vulnerabilties with HIGH or CRITICAL
    severity has been found.

    Args:
        repo (github.Repository): the repository where the issue will be posted.
        existing_issue (github.Issue): the duplicate on issue on GitHub for
            the current release, if it exists. Otherwise None.
        trivy_json (dict): parsed from Trivy's JSON result using load_results().
        trivy_table (str): Trivy results in a more readable table format which
            to be included in GitHub issues. Also read using load_results().

    Raises:
        SystemExit: If vulnerabilities with high or critical severity has been
            found.

    """
    vulnerabilities = trivy_json["Vulnerabilities"]
    num_vulnerabilities = len(vulnerabilities)
    severity_list = get_severity_list(vulnerabilities)

    # Create or update issue as required.
    title, body = build_issue(
        num_vulnerabilities, severity_list, trivy_table, trivy_json)
    issue_number = create_or_update_issue(
        repo, existing_issue, title, body)

    # Print details to build log.
    details = ("{} vulnerabilities of type: [{}] were found in image. "
               "Please refer to issue #{} for details.").format(
        num_vulnerabilities, ",".join(severity_list), issue_number)
    print(details)

    # Fail build if high or critical vulnerabilities are found.
    if high_or_critical_exists(severity_list):
        sys.exit("Build will be aborted due to high and/or critical "
                 "vulnerabilities found in the new images. The new "
                 "images will not be pushed to GCR.")
    else:
        print("Images will be published to GCR.")


def close_issue(existing_issue):
    """Closes existing issue on GitHub with a comment marking as resolved."""
    details = ("Closing Issue #{} since vulnerabilities have been resolved").format(
        existing_issue.number)
    print(details)

    existing_issue.create_comment(details)
    existing_issue.edit(state="closed")


def main(argv):
    """Parse the Trivy results."""
    validate_files()
    [trivy_json, trivy_table] = load_results()

    repo = get_github_repo()
    existing_issue = get_duplicate_issue(repo)

    if trivy_json["Vulnerabilities"]:
        process_vulnerabilities(repo, existing_issue, trivy_json, trivy_table)
    else:
        if existing_issue:
            close_issue(existing_issue)

        print("No vulnerabilities found. Images will be published to GCR.")


if __name__ == "__main__":
    flags.mark_flag_as_required("build_id")
    flags.mark_flag_as_required("commit_id")
    flags.mark_flag_as_required("release_tag")
    flags.mark_flag_as_required("repo_name")
    flags.mark_flag_as_required("token_file")
    flags.mark_flag_as_required("trivy_file")

    app.run(main)
