#!/usr/bin/env python3
import argparse
import json
import os
import subprocess
import sys
from pathlib import Path


def _run(cmd, check=True):
    return subprocess.run(cmd, check=check)


def _write_file(path: Path, content: str, mode: int = 0o644):
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(content, encoding="utf-8")
    os.chmod(tmp, mode)
    tmp.replace(path)


def _load_profile(profile_name: str) -> dict:
    base = Path("/usr/share/saikyo-security-profile/profiles")
    p = base / f"{profile_name}.json"
    if not p.exists():
        raise FileNotFoundError(f"profile not found: {p}")
    return json.loads(p.read_text(encoding="utf-8"))


def apply_profile(profile_name: str):
    prof = _load_profile(profile_name)

    sysctl_lines = prof.get("sysctl", {})
    sysctl_text = "".join([f"{k} = {v}\n" for k, v in sysctl_lines.items()])
    _write_file(Path("/etc/sysctl.d/99-saikyo-security-profile.conf"), sysctl_text)
    _run(["/sbin/sysctl", "--system"], check=False)

    enable_services = prof.get("enable_services", [])
    for svc in enable_services:
        _run(["/bin/systemctl", "enable", "--now", svc], check=False)

    ufw = prof.get("ufw", {})
    if ufw.get("enable", False):
        _run(["/usr/sbin/ufw", "--force", "reset"], check=False)
        default_in = ufw.get("default_in", "deny")
        default_out = ufw.get("default_out", "allow")
        _run(["/usr/sbin/ufw", "default", default_in, "incoming"], check=False)
        _run(["/usr/sbin/ufw", "default", default_out, "outgoing"], check=False)
        for rule in ufw.get("allow", []):
            _run(["/usr/sbin/ufw", "allow"] + rule.split(), check=False)
        _run(["/usr/sbin/ufw", "--force", "enable"], check=False)


def export_profile(profile_name: str, out_path: str):
    prof = _load_profile(profile_name)
    out = Path(out_path)
    out.write_text(json.dumps(prof, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def import_profile(profile_name: str, in_path: str):
    base = Path("/usr/share/saikyo-security-profile/profiles")
    dst = base / f"{profile_name}.json"
    src = Path(in_path)
    data = json.loads(src.read_text(encoding="utf-8"))
    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def main(argv):
    p = argparse.ArgumentParser(prog="saikyo-security-profile")
    sub = p.add_subparsers(dest="cmd", required=True)

    ap = sub.add_parser("apply")
    ap.add_argument("profile", choices=["standard", "secure"])

    ep = sub.add_parser("export")
    ep.add_argument("profile", choices=["standard", "secure"])
    ep.add_argument("path")

    ip = sub.add_parser("import")
    ip.add_argument("profile", choices=["standard", "secure"])
    ip.add_argument("path")

    args = p.parse_args(argv)

    if args.cmd == "apply":
        apply_profile(args.profile)
        return 0
    if args.cmd == "export":
        export_profile(args.profile, args.path)
        return 0
    if args.cmd == "import":
        import_profile(args.profile, args.path)
        return 0

    return 2


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
