#!/usr/bin/env python3

from __future__ import annotations

import argparse
import dataclasses
import hashlib
import json
import os
import random
import string
import sys
from typing import Iterable, List, Optional, Sequence

import requests

EXIT_OK = 101
EXIT_CORRUPT = 102
EXIT_MUMBLE = 103
EXIT_DOWN = 104
EXIT_CHECKER_ERROR = 110

SCHEME = os.getenv("SIGMA_SCHEME", "http")
PORT = int(os.getenv("SIGMA_API_PORT") or "3000")
API_PREFIX = os.getenv("SIGMA_API_PREFIX", "/api/v1").rstrip("/")
REQUEST_TIMEOUT = float(os.getenv("SIGMA_TIMEOUT", "5"))

ALPHABET = string.ascii_lowercase + string.digits
PASSWORD_ALPHABET = string.ascii_letters + string.digits


def log(message: str) -> None:
    print(message, file=sys.stderr)


class CheckerError(Exception):
    exit_code = EXIT_MUMBLE

    def __init__(self, message: str):
        super().__init__(message)


class DownError(CheckerError):
    exit_code = EXIT_DOWN


class CorruptError(CheckerError):
    exit_code = EXIT_CORRUPT


class CheckerInternalError(CheckerError):
    exit_code = EXIT_CHECKER_ERROR


class ApiClient:
    def __init__(self, base_url: str, timeout: float = REQUEST_TIMEOUT):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.session = requests.Session()
        self.session.headers.update({"User-Agent": "sigma-checker/1.0"})

    def request(
        self,
        method: str,
        path: str,
        *,
        expected: Optional[Iterable[int]] = None,
        **kwargs,
    ) -> requests.Response:
        url = self._build_url(path)
        try:
            response = self.session.request(
                method,
                url,
                timeout=self.timeout,
                **kwargs,
            )
        except requests.RequestException as exc:
            raise DownError(
                f"network error during {method.upper()} {path}: {exc}"
            ) from exc

        if expected is not None:
            allowed = set(expected)
            if response.status_code not in allowed:
                body = response.text[:400]
                raise CheckerError(
                    "unexpected status "
                    f"{response.status_code} for "
                    f"{method.upper()} {path}: {body}"
                )

        return response

    def _build_url(self, path: str) -> str:
        if not path.startswith("/"):
            path = "/" + path
        return f"{self.base_url}{path}"


def seeded_random(flag_id: str, vuln: int) -> random.Random:
    digest = hashlib.sha256(f"{flag_id}:{vuln}".encode()).digest()
    return random.Random(int.from_bytes(digest, "big"))


def random_string(
    rng: random.Random,
    length: int,
    alphabet: str = ALPHABET,
) -> str:
    return "".join(rng.choice(alphabet) for _ in range(length))


@dataclasses.dataclass
class UserSpec:
    username: str
    password: str


@dataclasses.dataclass
class ProjectSpec:
    name: str
    slug: str
    description_template: str


@dataclasses.dataclass
class FunctionSpec:
    name: str
    path: str


@dataclasses.dataclass
class Vuln1Context:
    user: UserSpec
    primary_project: ProjectSpec
    decoy_project: ProjectSpec


@dataclasses.dataclass
class Vuln2Context:
    user: UserSpec
    project: ProjectSpec
    store_function: FunctionSpec
    fetch_function: FunctionSpec


def build_vuln1_context(flag_id: str) -> Vuln1Context:
    rng = seeded_random(flag_id, 1)
    username = f"user_{random_string(rng, 10)}"
    password = random_string(rng, 16, PASSWORD_ALPHABET)
    primary = ProjectSpec(
        name=f"Sigma {random_string(rng, 5).title()}",
        slug=f"proj{random_string(rng, 10)}",
        description_template="Primary telemetry secret: {flag}",
    )
    decoy = ProjectSpec(
        name=f"Telemetry {random_string(rng, 5).title()}",
        slug=f"proj{random_string(rng, 10)}",
        description_template="Routine system baseline",
    )
    return Vuln1Context(UserSpec(username, password), primary, decoy)


def build_vuln2_context(flag_id: str) -> Vuln2Context:
    rng = seeded_random(flag_id, 2)
    username = f"{random_string(rng, 12)}"
    password = random_string(rng, 18, PASSWORD_ALPHABET)
    project = ProjectSpec(
        name=f"Runner {random_string(rng, 6).title()}",
        slug=f"proj{random_string(rng, 12)}",
        description_template="Serverless SQL vault",
    )
    store_fn = FunctionSpec(
        name=f"store_{random_string(rng, 6)}",
        path=f"/{random_string(rng, 6)}",
    )
    fetch_fn = FunctionSpec(
        name=f"fetch_{random_string(rng, 6)}",
        path=f"/{random_string(rng, 6)}",
    )
    return Vuln2Context(
        UserSpec(username, password),
        project,
        store_fn,
        fetch_fn,
    )


class SigmaChecker:
    def __init__(
        self,
        host: str,
        *,
        scheme: str = SCHEME,
        port: Optional[int] = PORT,
    ) -> None:
        base_url = self._build_base_url(host, scheme, port)
        prefix = API_PREFIX if API_PREFIX else ""
        self.client = ApiClient(f"{base_url}{prefix}")

    @staticmethod
    def _build_base_url(host: str, scheme: str, port: Optional[int]) -> str:
        hostname = host
        if ":" in host and not host.startswith("["):
            hostname = f"[{host}]"
        if port is None:
            return f"{scheme}://{hostname}"
        return f"{scheme}://{hostname}:{port}"

    def check(self) -> None:
        response = self.client.request("GET", "/health/cpu", expected={200})
        try:
            payload = response.json()
        except ValueError as exc:
            raise CheckerError(
                "health endpoint returned invalid JSON"
            ) from exc
        usage = payload.get("usage")
        if not isinstance(usage, (int, float)):
            raise CheckerError("health payload missing numeric usage")

    def put(self, flag_id: str, flag: str, vuln: int) -> None:
        if vuln == 1:
            self._put_vuln1(build_vuln1_context(flag_id), flag)
        elif vuln == 2:
            self._put_vuln2(build_vuln2_context(flag_id), flag)
        else:
            raise CheckerInternalError(f"unsupported vuln {vuln}")

    def get(self, flag_id: str, flag: str, vuln: int) -> None:
        if vuln == 1:
            self._get_vuln1(build_vuln1_context(flag_id), flag)
        elif vuln == 2:
            self._get_vuln2(build_vuln2_context(flag_id), flag)
        else:
            raise CheckerInternalError(f"unsupported vuln {vuln}")

    def _ensure_user(self, user: UserSpec) -> None:
        payload = {"username": user.username, "password": user.password}
        response = self.client.request(
            "POST",
            "/users",
            json=payload,
            expected={201, 409},
        )
        if response.status_code == 201:
            log(f"registered {user.username}")

    def _login(self, user: UserSpec, *, corrupt_on_fail: bool = False) -> str:
        payload = {"username": user.username, "password": user.password}
        response = self.client.request(
            "POST",
            "/auth/login",
            json=payload,
            expected={200, 401},
        )
        if response.status_code == 401:
            if corrupt_on_fail:
                raise CorruptError("stored credentials rejected")
            raise CheckerError("authentication failed")
        try:
            body = response.json()
        except ValueError as exc:
            raise CheckerError("login endpoint returned invalid JSON") from exc
        token = body.get("auth_token")
        if not isinstance(token, str) or not token:
            raise CheckerError("login response missing token")
        return token

    def _auth_headers(self, token: str) -> dict:
        return {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }

    def _create_project(
        self,
        headers: dict,
        spec: ProjectSpec,
        flag: Optional[str] = None,
    ) -> dict:
        description = spec.description_template
        if flag is not None:
            description = description.format(flag=flag)
        payload = {
            "name": spec.name,
            "slug": spec.slug,
            "description": description,
        }
        response = self.client.request(
            "POST",
            "/projects",
            headers=headers,
            json=payload,
            expected={201},
        )
        try:
            project = response.json()
        except ValueError as exc:
            raise CheckerError(
                "project creation returned invalid JSON"
            ) from exc
        if not isinstance(project.get("_id"), str):
            raise CheckerError("project creation missing id")
        return project

    def _list_projects(self, headers: dict) -> List[dict]:
        response = self.client.request(
            "GET",
            "/projects",
            headers=headers,
            expected={200},
        )
        try:
            projects = response.json()
        except ValueError as exc:
            raise CheckerError("projects list returned invalid JSON") from exc
        if not isinstance(projects, list):
            raise CheckerError("projects response is not a list")
        return projects

    def _find_project_by_slug(
        self,
        headers: dict,
        slug: str,
    ) -> Optional[dict]:
        for project in self._list_projects(headers):
            if project.get("slug") == slug:
                return project
        return None

    def _create_function(
        self,
        headers: dict,
        project_id: str,
        spec: FunctionSpec,
    ) -> dict:
        payload = {"name": spec.name, "path": spec.path}
        response = self.client.request(
            "POST",
            f"/projects/{project_id}/functions",
            headers=headers,
            json=payload,
            expected={201},
        )
        try:
            func = response.json()
        except ValueError as exc:
            raise CheckerError(
                "function creation returned invalid JSON"
            ) from exc
        if not isinstance(func.get("_id"), str):
            raise CheckerError("function creation missing id")
        return func

    def _update_function(
        self,
        headers: dict,
        project_id: str,
        function_id: str,
        code: str,
        methods: Sequence[str],
    ) -> None:
        payload = {"code": code, "methods": list(methods)}
        self.client.request(
            "PUT",
            f"/projects/{project_id}/functions/{function_id}",
            headers=headers,
            json=payload,
            expected={200},
        )

    def _execute_function(
        self,
        method: str,
        slug: str,
        path: str,
        expected: Iterable[int] = (200,),
    ) -> requests.Response:
        return self.client.request(
            method,
            f"/exec/{slug}{path}",
            expected=set(expected),
        )

    def _put_vuln1(self, context: Vuln1Context, flag: str) -> None:
        self._ensure_user(context.user)
        token = self._login(context.user)
        headers = self._auth_headers(token)
        self._create_project(headers, context.primary_project, flag)
        self._create_project(headers, context.decoy_project)
        stored = self._find_project_by_slug(
            headers,
            context.primary_project.slug,
        )
        if not stored:
            raise CheckerError("project missing after creation")
        description = stored.get("description", "")
        if flag not in description:
            raise CheckerError("flag not present immediately after storing")

    def _get_vuln1(self, context: Vuln1Context, flag: str) -> None:
        token = self._login(context.user, corrupt_on_fail=True)
        headers = self._auth_headers(token)
        project = self._find_project_by_slug(
            headers,
            context.primary_project.slug,
        )
        if not project:
            raise CorruptError("flag project missing")
        description = project.get("description", "")
        if flag not in description:
            raise CorruptError("flag missing from project description")

    def _put_vuln2(self, context: Vuln2Context, flag: str) -> None:
        self._ensure_user(context.user)
        token = self._login(context.user)
        headers = self._auth_headers(token)
        project = self._create_project(headers, context.project)
        project_id = project["_id"]
        func = self._create_function(
            headers,
            project_id,
            context.store_function,
        )
        code = self._build_store_function_code(context.project.slug, flag)
        self._update_function(headers, project_id, func["_id"], code, ["POST"])
        self._execute_function(
            "POST",
            context.project.slug,
            context.store_function.path,
            expected=(200, 201),
        )

    def _get_vuln2(self, context: Vuln2Context, flag: str) -> None:
        token = self._login(context.user, corrupt_on_fail=True)
        headers = self._auth_headers(token)
        project = self._find_project_by_slug(headers, context.project.slug)
        if not project:
            raise CorruptError("project not found during GET")
        project_id = project.get("_id")
        if not isinstance(project_id, str):
            raise CheckerError("project missing id")
        func = self._create_function(
            headers,
            project_id,
            context.fetch_function,
        )
        code = self._build_fetch_function_code(context.project.slug)
        self._update_function(headers, project_id, func["_id"], code, ["GET"])
        response = self._execute_function(
            "GET",
            context.project.slug,
            context.fetch_function.path,
            expected=(200, 404),
        )
        if response.status_code == 404:
            raise CorruptError("flag table returned empty result")
        payload = self._parse_json(response.text)
        if payload.get("flag") != flag:
            raise CorruptError("retrieved flag mismatch")

    def _build_store_function_code(self, slug: str, flag: str) -> str:
        create_sql = (
            f'CREATE TABLE IF NOT EXISTS "scope_{slug}" '
            "(id SERIAL PRIMARY KEY, flag TEXT NOT NULL)"
        )
        insert_sql = f'INSERT INTO "scope_{slug}" (flag) VALUES ($1)'
        return (
            """
;(async () => {
  const CREATE_SQL = __CREATE__;
  const INSERT_SQL = __INSERT__;
  const FLAG_VALUE = __FLAG__;
  await sql.execute(CREATE_SQL);
  await sql.execute(INSERT_SQL, [FLAG_VALUE]);
  res.setStatus(201).json({ status: "stored" });
})().catch((err) => {
  res.setStatus(500).json({ error: String(err?.message ?? err) });
});
"""
            .replace("__CREATE__", json.dumps(create_sql))
            .replace("__INSERT__", json.dumps(insert_sql))
            .replace("__FLAG__", json.dumps(flag))
            .strip()
        )

    def _build_fetch_function_code(self, slug: str) -> str:
        select_sql = f'SELECT flag FROM "scope_{slug}" ORDER BY id DESC LIMIT 1'
        return (
            """
;(async () => {
  const SELECT_SQL = __SELECT__;
  const rows = await sql.query(SELECT_SQL);
    if (
        !Array.isArray(rows) ||
        rows.length === 0 ||
        typeof rows[0]?.flag !== "string"
    ) {
    res.setStatus(404).json({ error: "flag missing" });
    return;
  }
  res.setStatus(200).json({ flag: rows[0].flag });
})().catch((err) => {
  res.setStatus(500).json({ error: String(err?.message ?? err) });
});
"""
            .replace("__SELECT__", json.dumps(select_sql))
            .strip()
        )

    @staticmethod
    def _parse_json(body: str) -> dict:
        try:
            data = json.loads(body or "null")
        except json.JSONDecodeError as exc:
            raise CheckerError("function response is not valid JSON") from exc
        if not isinstance(data, dict):
            raise CheckerError("function response JSON is not an object")
        return data


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Sigma ForcAD checker")
    parser.add_argument(
        "action",
        choices=["check", "put", "get"],
        help="Checker action",
    )
    parser.add_argument("host", help="Team IP address")
    parser.add_argument("flag_id", nargs="?")
    parser.add_argument("flag", nargs="?")
    parser.add_argument("vuln", nargs="?", type=int)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    checker = SigmaChecker(args.host)
    try:
        if args.action == "check":
            checker.check()
        elif args.action == "put":
            if not (args.flag_id and args.flag and args.vuln is not None):
                raise CheckerInternalError("put requires flag_id flag vuln")
            checker.put(args.flag_id, args.flag, args.vuln)
        elif args.action == "get":
            if not (args.flag_id and args.flag and args.vuln is not None):
                raise CheckerInternalError("get requires flag_id flag vuln")
            checker.get(args.flag_id, args.flag, args.vuln)
        else:
            raise CheckerInternalError("unknown action")
    except CheckerError as exc:
        log(str(exc))
        return exc.exit_code
    except Exception as exc:  # pragma: no cover
        log(f"unexpected checker error: {exc}")
        return EXIT_CHECKER_ERROR
    log("OK")
    return EXIT_OK


if __name__ == "__main__":
    sys.exit(main())