#!/usr/bin/env python3

import argparse
import json
import os
import random
import secrets
import string
import subprocess
import time
import traceback
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from enum import Enum
from pathlib import Path
from threading import Lock, current_thread
from typing import List, Tuple

import yaml
from dockerfile_parse import DockerfileParser

BASE_DIR = Path(__file__).resolve().absolute().parent
SERVICES_PATH = BASE_DIR / "services"
CHECKERS_PATH = BASE_DIR / "checkers"
MAX_THREADS = int(os.getenv("MAX_THREADS", default=2 * os.cpu_count()))
RUNS = int(os.getenv("RUNS", default=10))
HOST = os.getenv("HOST", default="127.0.0.1")
OUT_LOCK = Lock()
DISABLE_LOG = False

DC_REQUIRED_OPTIONS = ["services"]
DC_ALLOWED_OPTIONS = DC_REQUIRED_OPTIONS + ["volumes", "version"]

CONTAINER_REQUIRED_OPTIONS = ["restart"]
CONTAINER_ALLOWED_OPTIONS = CONTAINER_REQUIRED_OPTIONS + [
    "pids_limit",
    "mem_limit",
    "cpus",
    "build",
    "image",
    "ports",
    "volumes",
    "environment",
    "env_file",
    "healthcheck",
    "depends_on",
    "sysctls",
    "privileged",
    "security_opt",
]
SERVICE_REQUIRED_OPTIONS = ["pids_limit", "mem_limit", "cpus"]
SERVICE_ALLOWED_OPTIONS = CONTAINER_ALLOWED_OPTIONS
DATABASES = [
    "redis",
    "postgres",
    "mysql",
    "mariadb",
    "mongo",
    "mssql",
    "clickhouse",
    "tarantool",
]
PROXIES = ["nginx", "envoy"]
CLEANERS = ["dedcleaner"]

VALIDATE_DIRS = ["checkers", "services", "internal", "sploits"]

ALLOWED_CHECKER_PATTERNS = [
    "import requests",
    "requests.exceptions",
    "s: requests.Session",
    "sess: requests.Session",
    "session: requests.Session",
    "r: requests.Response",
    "resp: requests.Response",
    "Got requests connection error",
]
FORBIDDEN_CHECKER_PATTERNS = ["requests"]

ALLOWED_YAML_FILES = [
    "buf.yaml",
    "buf.gen.yaml",
    "application.yaml",
]


class ColorType(Enum):
    INFO = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    BOLD = "\033[1m"
    ENDC = "\033[0m"

    def __str__(self):
        return self.value


def generate_flag(name):
    alph = string.ascii_uppercase + string.digits
    return name[0].upper() + "".join(random.choices(alph, k=30)) + "="


def colored_log(*messages, color: ColorType = ColorType.INFO):
    ts = datetime.utcnow().isoformat(sep=" ", timespec="milliseconds")
    print(
        f"{color}{color.name} [{current_thread().name} {ts}]{ColorType.ENDC}", *messages
    )


class BaseValidator:
    def _log(self, message: str):
        with OUT_LOCK:
            if not DISABLE_LOG:
                colored_log(f"{self}: {message}")

    def _fatal(self, cond, message):
        global DISABLE_LOG

        with OUT_LOCK:
            if not cond:
                if not DISABLE_LOG:
                    colored_log(f"{self}: {message}", color=ColorType.FAIL)
                DISABLE_LOG = True
                raise AssertionError

    def _warning(self, cond: bool, message: str) -> bool:
        with OUT_LOCK:
            if not cond and not DISABLE_LOG:
                colored_log(f"{self}: {message}", color=ColorType.WARNING)
        return not cond

    def _error(self, cond, message) -> bool:
        with OUT_LOCK:
            if not cond and not DISABLE_LOG:
                colored_log(f"{self}: {message}", color=ColorType.FAIL)
        return not cond


class Checker(BaseValidator):
    def __init__(self, name: str):
        self._name = name
        self._exe_path = CHECKERS_PATH / self._name / "checker.py"
        self._fatal(
            os.access(self._exe_path, os.X_OK),
            f"{self._exe_path.relative_to(BASE_DIR)} must be executable",
        )
        self._timeout = 3
        self._get_info()

    def _get_info(self):
        self._log("running info action")
        cmd = [str(self._exe_path), "info", HOST]
        out, _ = self._run_command(cmd)
        info = json.loads(out)
        self._log(f"got info: {info}")

        self._vulns = int(info["vulns"])
        self._timeout = int(info["timeout"])
        self._attack_data = bool(info["attack_data"])

        self._fatal(
            60 > self._timeout > 0,
            f"invalid timeout: {self._timeout}",
        )

    @property
    def info(self):
        return {
            "vulns": self._vulns,
            "timeout": self._timeout,
            "attack_data": self._attack_data,
        }

    def _run_command(self, command: List[str], env=None) -> Tuple[str, str]:
        action = command[1].upper()
        cmd = ["timeout", str(self._timeout)] + command

        if env is None:
            env = os.environ
        env["PYTHONUNBUFFERED"] = "1"
        env["PWNLIB_NOTERM"] = "1"

        start = time.monotonic()
        p = subprocess.run(cmd, capture_output=True, check=False, env=env)
        elapsed = time.monotonic() - start

        out = p.stdout.decode()
        err = p.stderr.decode()

        out_s = out.rstrip("\n")
        err_s = err.rstrip("\n")

        self._log(
            f"action: {action}\ntime: {elapsed:.2f}s\nstdout:\n{out_s}\nstderr:\n{err_s}"
        )
        self._fatal(
            p.returncode != 124,
            f"action {action}: bad return code: 124, probably {ColorType.BOLD}timeout{ColorType.ENDC}",
        )
        self._fatal(
            p.returncode == 101, f"action {action}: bad return code: {p.returncode}"
        )
        return out, err

    def check(self):
        self._log("running CHECK")
        cmd = [str(self._exe_path), "check", HOST]
        self._run_command(cmd)

    def put(self, flag: str, flag_id: str, vuln: int):
        self._log(f"running PUT, flag={flag} flag_id={flag_id} vuln={vuln}")
        cmd = [str(self._exe_path), "put", HOST, flag_id, flag, str(vuln)]
        out, err = self._run_command(cmd)

        self._fatal(len(out) <= 1024, "returned stdout is longer than 1024 characters")
        self._fatal(len(err) <= 1024, "returned stderr is longer than 1024 characters")

        if self._attack_data:
            self._fatal(out, "stdout is empty")
            self._fatal(err, "stderr is empty")

            self._fatal(flag not in out, "flag is leaked in public data")

            # new flag ID is in stderr for attack_data checkers
            return err

        self._fatal(out, "stdout is empty")

        # new flag ID is in stdout for checkers without attack_data
        return out

    def get(self, flag: str, flag_id: str, vuln: int):
        self._log(f"running GET, flag={flag} flag_id={flag_id} vuln={vuln}")
        cmd = [str(self._exe_path), "get", HOST, flag_id, flag, str(vuln)]
        self._run_command(cmd)

    def run_all(self, step: int):
        self._log(f"running all actions (run {step} of {RUNS})")
        self.check()

        for vuln in range(1, self._vulns + 1):
            flag = generate_flag(self._name)
            flag_id = self.put(flag=flag, flag_id=secrets.token_hex(16), vuln=vuln)
            flag_id = flag_id.strip()
            self.get(flag, flag_id, vuln)

    def __str__(self):
        return f"checker {self._name}"


class Service(BaseValidator):
    def __init__(self, name: str):
        self._name = name
        self._path = SERVICES_PATH / self._name
        self._dc_path = self._path / "docker-compose.yml"
        self._fatal(
            self._dc_path.exists(),
            f"{self._dc_path.relative_to(BASE_DIR)} missing",
        )

        self._checker = Checker(self._name)

    @property
    def name(self):
        return self._name

    @property
    def checker_info(self):
        return self._checker.info

    def _run_dc(self, *args):
        cmd = ["docker", "compose", "-f", str(self._dc_path)] + list(args)
        subprocess.run(cmd, check=True)

    def up(self):
        self._log("starting")
        self._run_dc("up", "--build", "-d")

    def logs(self):
        self._log("printing logs")
        self._run_dc("logs", "--tail", "2000")

    def down(self):
        self._log("stopping")
        self._run_dc("down", "-v")

    def validate_checker(self):
        self._log("validating checker")

        cnt_threads = max(1, min(MAX_THREADS, RUNS // 10))
        self._log(f"starting {cnt_threads} checker threads")
        with ThreadPoolExecutor(
            max_workers=cnt_threads,
            thread_name_prefix="Executor",
        ) as executor:
            for _ in executor.map(self._checker.run_all, range(1, RUNS + 1)):
                pass

    def __str__(self):
        return f"service {self._name}"


class StructureValidator(BaseValidator):
    def __init__(self, d: Path, service: Service):
        self._dir = d
        self._was_error = False
        self._service = service

    def _error(self, cond, message):
        err = super()._error(cond, message)
        self._was_error |= err
        return err

    def validate(self):
        for d in VALIDATE_DIRS:
            self.validate_dir(self._dir / d / self._service.name)
        return not self._was_error

    def validate_dir(self, d: Path):
        if not d.exists():
            return
        for f in d.iterdir():
            if f.is_file():
                self.validate_file(f)
            elif f.name[0] != ".":
                self.validate_dir(f)

    def validate_file(self, f: Path):
        path = f.relative_to(BASE_DIR)

        if f.name not in ALLOWED_YAML_FILES:
            self._error(f.suffix != ".yaml", f"file {path} has .yaml extension")

        self._error(f.name != ".gitkeep", f"{path} found, should be named .keep")

        if f.name == "docker-compose.yml":
            with f.open() as file:
                dc = yaml.safe_load(file)

            if self._error(isinstance(dc, dict), f"{path} is not dict"):
                return

            for opt in DC_REQUIRED_OPTIONS:
                if self._error(opt in dc, f"required option {opt} not in {path}"):
                    return

            if "version" in dc:
                if self._error(
                    isinstance(dc["version"], str),
                    f"version option in {path} is not string",
                ):
                    return

                try:
                    dc_version = float(dc["version"])
                except ValueError:
                    self._error(False, f"version option in {path} is not float")
                    return

                self._error(
                    2.4 <= dc_version < 3,
                    f"invalid version in {path}, need >=2.4 and <3 (or no version at all), got {dc_version}",
                )

            for opt in dc:
                self._error(
                    opt in DC_ALLOWED_OPTIONS,
                    f"option {opt} in {path} is not allowed",
                )

            services = []
            databases = []
            proxies = []
            dependencies = defaultdict(list)

            if self._error(
                isinstance(dc["services"], dict),
                f"services option in {path} is not dict",
            ):
                return

            for container, container_conf in dc["services"].items():
                if self._error(
                    isinstance(container_conf, dict),
                    f"config in {path} for container {container} is not dict",
                ):
                    continue

                for opt in CONTAINER_REQUIRED_OPTIONS:
                    self._error(
                        opt in container_conf,
                        f"required option {opt} not in {path} for container {container}",
                    )

                self._error(
                    "restart" in container_conf
                    and container_conf["restart"] == "unless-stopped",
                    f'restart option in {path} for container {container} must be equal to "unless-stopped"',
                )

                for opt in container_conf:
                    self._error(
                        opt in CONTAINER_ALLOWED_OPTIONS,
                        f"option {opt} in {path} is not allowed for container {container}",
                    )

                if self._error(
                    "image" not in container_conf or "build" not in container_conf,
                    f"both image and build options in {path} for container {container}",
                ):
                    continue

                if self._error(
                    "image" in container_conf or "build" in container_conf,
                    f"both image and build options not in {path} for container {container}",
                ):
                    continue

                if "image" in container_conf:
                    image = container_conf["image"]
                else:
                    build = container_conf["build"]
                    if isinstance(build, str):
                        dockerfile = f.parent / build / "Dockerfile"
                    else:
                        context = build["context"]
                        if "dockerfile" in build:
                            dockerfile = f.parent / context / build["dockerfile"]
                        else:
                            dockerfile = f.parent / context / "Dockerfile"

                    if self._error(
                        dockerfile.exists(), f"no dockerfile found in {dockerfile}"
                    ):
                        continue

                    with dockerfile.open() as file:
                        dfp = DockerfileParser(fileobj=file)
                        image = dfp.baseimage

                    if self._error(
                        image is not None, f"no image option in {dockerfile}"
                    ):
                        continue

                if "depends_on" in container_conf:
                    for dependency in container_conf["depends_on"]:
                        dependencies[container].append(dependency)

                is_service = True
                for database in DATABASES:
                    if database in image:
                        databases.append(container)
                        is_service = False

                for proxy in PROXIES:
                    if proxy in image:
                        proxies.append(container)
                        is_service = False

                for cleaner in CLEANERS:
                    if cleaner in image:
                        is_service = False

                if is_service:
                    services.append(container)
                    for opt in SERVICE_REQUIRED_OPTIONS:
                        self._error(
                            opt in container_conf,
                            f"required option {opt} not in {path} for service {container}",
                        )

                    for opt in container_conf:
                        self._error(
                            opt in SERVICE_ALLOWED_OPTIONS,
                            f"option {opt} in {path} is not allowed for service {container}",
                        )

            for service in services:
                for database in databases:
                    self._warning(
                        service in dependencies and database in dependencies[service],
                        f"service {service} may need to depends_on database {database}",
                    )

            for proxy in proxies:
                for service in services:
                    self._warning(
                        proxy in dependencies and service in dependencies[proxy],
                        f"proxy {proxy} may need to depends_on service {service}",
                    )

        elif BASE_DIR / "checkers" in f.parents and f.suffix == ".py":
            checker_code = f.read_text()
            for p in ALLOWED_CHECKER_PATTERNS:
                checker_code = checker_code.replace(p, "")
            for p in FORBIDDEN_CHECKER_PATTERNS:
                self._error(p not in checker_code, f'forbidden pattern "{p}" in {path}')

    def __str__(self):
        return f"Structure validator for {self._service.name}"


def get_services() -> List[Service]:
    if os.getenv("SERVICE") in ["all", None]:
        result = list(
            Service(service_path.name)
            for service_path in SERVICES_PATH.iterdir()
            if service_path.name[0] != "." and service_path.is_dir()
        )
    else:
        result = [Service(os.environ["SERVICE"])]

    with OUT_LOCK:
        colored_log("Got services:", ", ".join(map(str, result)))
    return result


def list_services(_args):
    services = get_services()
    if outfile := os.getenv("GITHUB_OUTPUT"):
        data = {
            "include": [{"service": service.name} for service in services],
        }
        with open(outfile, "a") as f:
            f.write(f"matrix={json.dumps(data)}")


def start_services(_args):
    for service in get_services():
        service.up()


def stop_services(_args):
    for service in get_services():
        service.down()


def logs_services(_args):
    for service in get_services():
        service.logs()


def validate_checkers(_args):
    for service in get_services():
        service.validate_checker()


def validate_structure(_args):
    was_error = False
    for service in get_services():
        validator = StructureValidator(BASE_DIR, service)
        if not validator.validate():
            was_error = True

    if was_error:
        with OUT_LOCK:
            colored_log("Structure validator: failed", color=ColorType.FAIL)
            raise AssertionError


def dump_tasks(_args):
    result = {"tasks": []}
    for service in get_services():
        info = service.checker_info
        checker_type = "gevent"
        if info["attack_data"]:
            checker_type += "_pfr"

        result["tasks"].append(
            {
                "name": service.name,
                "checker": f"{service.name}/checker.py",
                "checker_timeout": info["timeout"],
                "checker_type": checker_type,
                "places": info["vulns"],
                "puts": 1,
                "gets": 1,
            }
        )

    colored_log("\n" + yaml.safe_dump(result))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Validate checkers for A&D. "
                    "Host & number of runs are passed with HOST and RUNS env vars"
    )
    subparsers = parser.add_subparsers()

    list_parser = subparsers.add_parser(
        "list",
        help="List services to test",
    )
    list_parser.set_defaults(func=list_services)

    up_parser = subparsers.add_parser(
        "up",
        help="Start services",
    )
    up_parser.set_defaults(func=start_services)

    down_parser = subparsers.add_parser(
        "down",
        help="Stop services",
    )
    down_parser.set_defaults(func=stop_services)

    logs_parser = subparsers.add_parser(
        "logs",
        help="Print logs for services",
    )
    logs_parser.set_defaults(func=logs_services)

    check_parser = subparsers.add_parser(
        "check",
        help="Run checkers validation",
    )
    check_parser.set_defaults(func=validate_checkers)

    validate_parser = subparsers.add_parser(
        "validate",
        help="Run structure validation",
    )
    validate_parser.set_defaults(func=validate_structure)

    dump_parser = subparsers.add_parser(
        "dump_tasks",
        help="Dump tasks in YAML for ForcAD",
    )
    dump_parser.set_defaults(func=dump_tasks)

    parsed = parser.parse_args()

    if "func" not in parsed:
        print("Type -h")
        exit(1)

    try:
        parsed.func(parsed)
    except AssertionError:
        exit(1)
    except Exception as e:
        tb = traceback.format_exc()
        print("Got exception, report it:", e, tb)
        exit(1)