#!/usr/bin/env python3

"""
check_container_scanning_job -- Monitoring plugin to check if a Docker image
scanned in a gitlab job has security vulnerabilities."

Author: Raphaël Laguerre <rlaguerre@easter-eggs.com>
Source: https://gitlab.easter-eggs.com/rlaguerre/check_container_scanning_job

Copyright (C) 2024 Easter-eggs

check_container_scanning_job is free software; you can redistribute
it and/or modify it under the terms of the GNU General Public License
as published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This software is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""


import argparse
import json
import logging
import os
import shutil
import sys
import urllib.parse
import zipfile
from datetime import datetime, timedelta
from enum import IntEnum
from getpass import getpass

import requests


class ExitCode(IntEnum):
    """Enum for exit codes"""

    OK = 0
    WARNING = 1
    CRITICAL = 2
    UNKNOWN = 3


MULTILINE = ""
STATUS = "OK"
MESSAGE = "No security vulnerabilities have been found"
CODE = ExitCode.OK


class Severity(IntEnum):
    """Trivy severity"""

    UNKNOWN = 0
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4


ZIP_FILE = "artifact.zip"
JSON_FILE = "gl-container-scanning-report.json"


def clean_exit():
    """Clean exit"""

    # Remove files
    shutil.rmtree(directory)
    # Exit
    doc = """Documentation can be found at https://www.easter-eggs.fr/ee/gitlab/accueil\
?s[]=scanning#job_de_scan_de_vulnerabilites_d_images_docker"""
    output = f"{STATUS}: {MESSAGE}" + "\n" + MULTILINE + doc
    print(output)
    sys.exit(CODE)


# Parse arguments
parser = argparse.ArgumentParser(
    description=""" Monitoring plugin to check if a Docker image scanned in a gitlab job has
        security vulnerabilities."""
)

log_args = parser.add_argument_group("Logging options")
log_args.add_argument("-d", "--debug", action="store_true", help="Show debug messages")

log_args.add_argument("-v", "--verbose", action="store_true", help="Show verbose messages")

log_args.add_argument("-w", "--warning", action="store_true", help="Show warning messages")

log_args.add_argument("-l", "--logfile", action="store", type=str, help="Log file path")

log_args.add_argument(
    "-c",
    "--console",
    action="store_true",
    help="Also log on console (even if log file is provided)",
)

gitlab_args = parser.add_argument_group("Gitlab options")
gitlab_args.add_argument(
    "-u",
    "--url",
    required=True,
    help="Gitlab URL",
)
gitlab_args.add_argument(
    "-t",
    "--token",
    help="Gitlab access token",
)
gitlab_args.add_argument(
    "-f",
    "--token-file",
    help="Gitlab access token file",
)
gitlab_args.add_argument(
    "-p",
    "--project",
    required=True,
    help="Gitlab project",
)
gitlab_args.add_argument(
    "-r", "--ref_name", help="Branch or tag name in repository", default="main"
)
gitlab_args.add_argument("-j", "--job", help="job name", default="container_scanning")
gitlab_args.add_argument(
    "-x",
    "--https-proxy",
    help="https proxy",
)

gitlab_args.add_argument(
    "-o",
    "--limit-old",
    help="If the job is older than this number of days, returns a critical status.",
    default="7",
    type=int,
)

gitlab_args.add_argument(
    "-C",
    "--limit-critical",
    help="If severities higher than limit-critical are found, return a critical status.",
    default="CRITICAL",
    choices=["UNKNOWN", "LOW", "MEDIUM", "HIGH", "CRITICAL"],
)

gitlab_args.add_argument(
    "-W",
    "--limit-warning",
    help="If severities higher than limit-warning are found, return a warning status.",
    default="HIGH",
    choices=["UNKNOWN", "LOW", "MEDIUM", "HIGH", "CRITICAL"],
)

args = parser.parse_args()

# Initialize log
log = logging.getLogger()
logformat = logging.Formatter(
    "%(asctime)s - " + os.path.basename(sys.argv[0]) + " - %(levelname)s : %(message)s"
)

if args.debug:
    log.setLevel(logging.DEBUG)
elif args.verbose:
    log.setLevel(logging.INFO)
elif args.warning:
    log.setLevel(logging.WARNING)
else:
    log.setLevel(logging.FATAL)

if args.logfile:
    logfile = logging.FileHandler(args.logfile)
    logfile.setFormatter(logformat)
    log.addHandler(logfile)

if not args.logfile or args.console:
    logconsole = logging.StreamHandler()
    logconsole.setFormatter(logformat)
    log.addHandler(logconsole)

# Enable warnings capture by Python logging to avoid displaying urllib warnings
# if logging is not enabled
logging.captureWarnings(True)

if not args.token:
    if args.token_file and os.path.exists(args.token_file):
        try:
            with open(args.token_file, encoding="utf8") as fd:
                content = fd.read().strip()
            args.token = content
        except Exception as err:  # pylint: disable=broad-exception-caught
            parser.exit(
                ExitCode.UNKNOWN,
                "UNKNOWN - Fail to load Gitlab access token from file "
                f"'{args.token_file}': {err}"
                "\n",
            )
        if not args.token:
            parser.exit(
                ExitCode.UNKNOWN,
                "UNKNOWN - Empty access token loaded from file "
                f" ({args.token_file}). Please check it or provide it "
                "using -t/--token parameter.\n",
            )
        else:
            log.debug("Gitlab token loaded from file '%s'", args.token_file)
    else:
        args.token = getpass(f"Please enter token for project {args.project}: ")

url = args.url
ref_name = args.ref_name
project = args.project
url_encoded_project = urllib.parse.quote(project, safe="")
url_encoded_ref_name = urllib.parse.quote(args.ref_name, safe="")
headers = {"PRIVATE-TOKEN": args.token}
proxy = {"https": args.https_proxy}
FULL_URL = (
    f"{url}/api/v4/projects/{url_encoded_project}/jobs/artifacts/{url_encoded_ref_name}/download"
)
params = {"job": args.job}

directory = f"""/var/lib/nagios/{url_encoded_project}-{url_encoded_ref_name}-
{datetime.now().strftime("%Y-%m-%dT%H:%M:%S:%f")}"""
ZIP_PATH = f"{directory}/{ZIP_FILE}"
JSON_PATH = f"{directory}/{JSON_FILE}"
limit_critical = Severity[args.limit_critical]
limit_warning = Severity[args.limit_warning]

try:
    os.mkdir(directory)
except Exception as err:  # pylint: disable=broad-exception-caught
    STATUS = "UNKNOWN"
    MESSAGE = err
    CODE = ExitCode.UNKNOWN
    clean_exit()

# Download the artifacts zipped archive from the latest successful pipeline for the given reference
# name and job, provided the job finished successfully.
try:
    r = requests.get(
        FULL_URL, headers=headers, params=params, proxies=proxy, stream=True, timeout=30
    )
    r.raise_for_status()
except Exception as err:  # pylint: disable=broad-exception-caught
    STATUS = "UNKNOWN"
    MESSAGE = err
    CODE = ExitCode.UNKNOWN
    clean_exit()

try:
    with open(ZIP_PATH, "wb") as fd:
        for chunk in r.iter_content(chunk_size=128):
            fd.write(chunk)
except Exception as err:  # pylint: disable=broad-exception-caught
    STATUS = "UNKNOWN"
    MESSAGE = f"""Error writing the zipped archive of the artifact of the last successful {args.job}
    job with the reference {ref_name} in the {project} project. {err}"""
    CODE = ExitCode.UNKNOWN
    clean_exit()

# Extract zipped report
try:
    with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
        zip_ref.extract(JSON_FILE, path=directory)
except Exception as err:  # pylint: disable=broad-exception-caught
    STATUS = "UNKNOWN"
    MESSAGE = f""""Error extracting the zipped archive of the artifact of the last successful
    {args.job} job with the reference {ref_name} in the {project} project. {err}"""
    CODE = ExitCode.UNKNOWN
    clean_exit()

# Read json report
try:
    with open(JSON_PATH, encoding="utf8") as fd:
        j = json.load(fd)
except Exception as err:  # pylint: disable=broad-exception-caught
    STATUS = "UNKNOWN"
    MESSAGE = f"""Error reading the json vulnerability scanning report of the artifact of the last
    successful {args.job} job with the reference {ref_name} in the {project} project. {err}"""
    CODE = ExitCode.UNKNOWN
    clean_exit()

# Evaluate result
str_job_date = j["scan"]["end_time"]
job_date = datetime.strptime(str_job_date, "%Y-%m-%dT%H:%M:%S")
limit_date = datetime.today() - timedelta(days=args.limit_old)
if job_date < limit_date:
    STATUS = "CRITICAL"
    MESSAGE = f"""Job {args.job} with the reference {ref_name} in the {project} project has been
    executed for the last time on {str_job_date} which is older than {args.limit_old} days ago."""
    CODE = ExitCode.CRITICAL
elif len(j["vulnerabilities"]) != 0:
    vulnerabilities = [0, 0, 0, 0, 0]
    for v in j["vulnerabilities"]:
        vulnerabilities[Severity[v.get("severity").upper()]] += 1
    for i in range(limit_warning, limit_critical):
        if vulnerabilities[Severity(i)] > 0:
            STATUS = "WARNING"
            CODE = ExitCode.WARNING
            break
    for i in range(limit_critical, len(Severity)):
        if vulnerabilities[Severity(i)] > 0:
            STATUS = "CRITICAL"
            CODE = ExitCode.CRITICAL
            break
    MESSAGE = f"""{len(j["vulnerabilities"])} security vulnerabilities have been found. Find a \
full report in {url}/{project}/-/artifacts"""
    MULTILINE += f"Number of critical vulnerabilities: {vulnerabilities[Severity.CRITICAL]}\n"
    MULTILINE += f"Number of high vulnerabilities: {vulnerabilities[Severity.HIGH]}\n"
    MULTILINE += f"Number of medium vulnerabilities: {vulnerabilities[Severity.MEDIUM]}\n"
    MULTILINE += f"Number of low vulnerabilities: {vulnerabilities[Severity.LOW]}\n"
    MULTILINE += f"Number of unknown vulnerabilities: {vulnerabilities[Severity.UNKNOWN]}\n"

clean_exit()

# vim: tabstop=4 shiftwidth=4 softtabstop=4 expandtab
