#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "ctfbridge",
#   "pyyaml",
#   "typer",
#   "rich",
# ]
# ///

"""
ctf-sniper: Automated flag submission helper for CTFs.

Usage:
    ./ctf-sniper.py config.yaml

License:
    MIT License (c) 2025 bjornmorten
"""

import asyncio
from pathlib import Path

import ctfbridge
import typer
import yaml
from pydantic import BaseModel, Field, ValidationError, model_validator
from rich.console import Console

app = typer.Typer(help="ctf-sniper")
console = Console()


# -------------------------
# Models
# -------------------------
class Auth(BaseModel):
    token: str | None = None
    username: str | None = None
    password: str | None = None

    @model_validator(mode="after")
    def check_auth(self):
        if self.token:
            if self.username or self.password:
                raise ValueError("Provide either token OR username+password, not both")
        else:
            if not (self.username and self.password):
                raise ValueError("Must provide either token OR username+password")
        return self


class StartConfig(BaseModel):
    start_time: str = None
    lead_ms: int = 200
    prewarm_seconds: float = 2.0


class Target(BaseModel):
    url: str
    auth: Auth


class FlagEntry(BaseModel):
    challenge_id: str | None = None
    match_terms: list[str] = Field(default_factory=list)
    flags: list[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def check_target(self):
        if not self.challenge_id and not self.match_terms:
            raise ValueError("Either challenge_id or match_terms must be provided")
        return self


class RootConfig(BaseModel):
    target: Target
    flags: list[FlagEntry]


# -------------------------
# Functions
# -------------------------
async def get_authenticated_client(target: Target):
    with console.status("[cyan]Connecting...[/cyan]") as status:
        try:
            client = await ctfbridge.create_client(target.url)
        except ctfbridge.exceptions.UnknownPlatformError:
            console.print("The CTF platform is not supported or could not be identified", style="red")
            raise typer.Exit(code=1)
        console.print(f"Connected to {target.url}", style="green")

        if not client.capabilities.submit_flags:
            console.print(
                f"Flag submission is not supported for the {client.platform_name} plaform",
                style="red",
            )
            raise typer.Exit(code=1)

        status.update("[cyan]Authenticating...[/cyan]")
        if target.auth.token:
            await client.auth.login(token=target.auth.token)
        else:
            await client.auth.login(username=target.auth.username, password=target.auth.password)

        console.print("Authenticated successfully", style="green")

    console.print("Client ready to submit flags", style="green")

    return client


async def snipe(config: RootConfig):
    client = await get_authenticated_client(config.target)

    # TODO: wait for the CTF to start

    # Start fetching challenges in the background if any lookup is needed
    needs_lookup = any(e.match_terms for e in config.flags)
    challenge_task = (
        asyncio.create_task(client.challenges.get_all(solved=False, detailed=True)) if needs_lookup else None
    )

    async def submit_for_entry(entry: FlagEntry):
        async def submit_flags(challenge_id: str) -> bool:
            for flag in entry.flags:
                try:
                    result = await client.challenges.submit(challenge_id, flag)
                    if result.correct:
                        console.print(f"Submitted {flag} → {result.message}", style="green")
                        return True
                    else:
                        console.print(f"Submitted {flag} → {result.message}", style="yellow")
                except Exception as e:
                    console.print(f"Error submitting {flag}: {e}", style="red")
            return False

        # 1. Try challenge id directly
        if entry.challenge_id:
            try:
                if await submit_flags(entry.challenge_id):
                    return
            except Exception as e:
                console.print(f"Failed with challenge_id {entry.challenge_id}, falling back: {e}", style="yellow")

        # 2. If fallback is possible, wait for challenges to be fetched
        if entry.match_terms:
            if not challenge_task:
                return
            challenges = await challenge_task

            def find_challenges():
                terms = [t.lower() for t in entry.match_terms]
                results = []
                for c in challenges:
                    searchable = " ".join(
                        [
                            c.name,
                            c.description,
                            c.category,
                        ]
                    ).lower()
                    if any(term in searchable for term in terms):
                        results.append(c)
                return results

            matches = find_challenges()
            if not matches:
                console.print(f"No challenge found with terms {entry.match_terms}", style="yellow")
                return

            for challenge in matches:
                success = await submit_flags(challenge.id)
                if success:
                    console.print(f"Solved {challenge.name} with flags from entry {entry}", style="green")

    await asyncio.gather(*(submit_for_entry(e) for e in config.flags))


# -------------------------
# CLI
# -------------------------
@app.command()
def run(
    config_file: Path = typer.Argument(
        ...,
        exists=True,
        file_okay=True,
        dir_okay=False,
        readable=True,
        resolve_path=True,
    ),
):
    try:
        data = yaml.safe_load(config_file.read_text(encoding="utf-8")) or {}
        config = RootConfig(**data)
    except yaml.YAMLError as e:
        console.print(f"YAML error in {config_file}: {e}", style="red")
        raise typer.Exit(1)
    except ValidationError as e:
        console.print("Config validation failed:", style="red")
        console.print(e)
        raise typer.Exit(2)

    asyncio.run(snipe(config))


def main():
    app()


if __name__ == "__main__":
    main()
