#!/usr/bin/env python3
"""
sftp_debug.py — Step-by-step SFTP connection debugger.

Modes:
  openssh  (default)
    1. TCP reachability check
    2. ssh-keyscan to retrieve host identity; verify against known_hosts or a
       supplied expected key; write result to a temp known_hosts file
    3. sftp -vvv to enumerate host key algorithms offered by the server
    4. sftp -vvv with private key to attempt authentication

  paramiko
    1. TCP reachability check
    2. Fetch and verify host key via paramiko Transport
    3. Authenticate and open an SFTP session via paramiko SSHClient
"""

import argparse
import base64
import os
import posixpath
import re
import socket
import subprocess
import sys
import tempfile
from pathlib import Path


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def fail(message: str) -> None:
    print(f"\n[FAIL] {message}", file=sys.stderr)
    sys.exit(1)


def section(title: str) -> None:
    bar = "=" * 60
    print(f"\n{bar}")
    print(f"  {title}")
    print(bar)


def run(cmd: list[str], *, timeout: int = 30, input: str | None = None) -> subprocess.CompletedProcess:
    """Run a command, capturing stdout and stderr."""
    return subprocess.run(
        cmd,
        capture_output=True,
        text=True,
        timeout=timeout,
        input=input,
    )


def _normalize_key(raw: str) -> str:
    """
    Reduce a key string to 'key-type base64' for comparison.
    Handles both bare keys ('ssh-ed25519 AAAA...') and known_hosts lines
    ('host ssh-ed25519 AAAA...' or '|1|hash|hash ssh-ed25519 AAAA...').
    """
    parts = raw.strip().split()
    # known_hosts line: hostname keytype base64 [comment]
    # bare key:         keytype base64 [comment]
    # Detect by checking whether parts[0] looks like a key type
    key_types = {"ssh-rsa", "ssh-dss", "ssh-ed25519", "ecdsa-sha2-nistp256",
                 "ecdsa-sha2-nistp384", "ecdsa-sha2-nistp521", "ssh-ed448"}
    if parts and parts[0] in key_types:
        return f"{parts[0]} {parts[1]}"
    if len(parts) >= 3 and parts[1] in key_types:
        return f"{parts[1]} {parts[2]}"
    return raw.strip()


# ---------------------------------------------------------------------------
# Step 1: TCP reachability
# ---------------------------------------------------------------------------

def check_reachability(host: str, port: int) -> None:
    section(f"Step 1: TCP reachability — {host}:{port}")
    try:
        with socket.create_connection((host, port), timeout=10):
            print(f"[OK] Reached {host}:{port}")
    except OSError as exc:
        fail(f"Cannot reach {host}:{port} — {exc}")


# ---------------------------------------------------------------------------
# Step 2: ssh-keyscan + host key verification
# ---------------------------------------------------------------------------

def _check_against_expected(scanned_keys: list[str], expected_key: str) -> None:
    """Fail if none of the scanned keys match the expected key."""
    expected_norm = _normalize_key(expected_key)
    for key_line in scanned_keys:
        if _normalize_key(key_line) == expected_norm:
            print(f"[OK] Scanned key matches provided expected key.")
            return
    print(f"[WARN] Expected key : {expected_norm}")
    print(f"[WARN] Scanned keys : {[_normalize_key(k) for k in scanned_keys]}")
    fail("Host key returned by server does not match --expected-host-key.")


def _check_against_known_hosts(host: str, port: int, scanned_keys: list[str]) -> None:
    """
    Compare scanned keys against ~/.ssh/known_hosts using ssh-keygen -F.
    Warns on mismatch but only fails if an entry exists and conflicts.
    """
    lookup_host = f"[{host}]:{port}" if port != 22 else host
    result = run(["ssh-keygen", "-F", lookup_host], timeout=10)

    if result.returncode != 0 or not result.stdout.strip():
        print("[INFO] No existing entry in ~/.ssh/known_hosts for this host.")
        return

    existing_lines = [l for l in result.stdout.splitlines() if not l.startswith("#")]
    if not existing_lines:
        print("[INFO] No existing entry in ~/.ssh/known_hosts for this host.")
        return

    print(f"[INFO] Existing known_hosts entry found:\n       {existing_lines[0].strip()}")

    existing_norm = {_normalize_key(l) for l in existing_lines}
    scanned_norm  = {_normalize_key(k) for k in scanned_keys}

    if existing_norm & scanned_norm:
        print("[OK] Scanned key matches existing known_hosts entry.")
    else:
        print("[WARN] Existing known_hosts entry : " + ", ".join(existing_norm))
        print("[WARN] Key returned by server     : " + ", ".join(scanned_norm))
        fail(
            "Host key mismatch — the server's key differs from your known_hosts entry.\n"
            "       If the server key legitimately changed, remove the old entry with:\n"
            f"       ssh-keygen -R {lookup_host}"
        )


def keyscan(host: str, port: int, expected_key: str | None) -> Path:
    """
    Run ssh-keyscan, verify the result, and return a Path to a temp
    known_hosts file populated with the scanned keys.
    """
    section(f"Step 2: Host key verification — {host}:{port}")
    cmd = ["ssh-keyscan", "-p", str(port), host]
    print(f"$ {' '.join(cmd)}\n")

    result = run(cmd, timeout=15)

    if result.stderr:
        print(result.stderr.rstrip())

    if not result.stdout.strip():
        fail(
            "ssh-keyscan returned no keys. Possible causes:\n"
            "  - The server is not running SSH/SFTP on this port\n"
            "  - Something else is listening on the port (proxy, load balancer)\n"
            "  - The server closed the connection before completing key exchange\n"
            "    (e.g. MaxStartups limit reached, or it's still starting up)\n"
            "  - The server uses a non-standard key exchange this tool doesn't probe\n"
            "  - A slow server response exceeded ssh-keyscan's internal timeout\n"
            "  - DNS resolved successfully but points to the wrong host"
        )

    print(result.stdout.rstrip())

    # Parse out the key lines (skip comments)
    scanned_keys = [l for l in result.stdout.splitlines() if l and not l.startswith("#")]

    if expected_key:
        _check_against_expected(scanned_keys, expected_key)
    else:
        _check_against_known_hosts(host, port, scanned_keys)

    # Write scanned keys to a temp known_hosts file for subsequent steps
    tmp = tempfile.NamedTemporaryFile(
        mode="w", prefix="sftp_debug_known_hosts_", delete=False, suffix=".txt"
    )
    tmp.write(result.stdout)
    tmp.close()
    print(f"\n[OK] Scanned keys written to temp known_hosts: {tmp.name}")

    return Path(tmp.name)


# ---------------------------------------------------------------------------
# Step 3: enumerate host key algorithms via sftp -vvv
# ---------------------------------------------------------------------------

_HOST_KEY_ALGO_PATTERN = re.compile(r"debug\d+:\s+host key algorithms:\s*(.+)", re.IGNORECASE)
_SERVER_HOST_KEY_ALGO_PATTERNS = [
    # OpenSSH client: server_host_key_algorithms is sometimes on its own line
    re.compile(r"server_host_key_algorithms[=:\s]+(.+)", re.IGNORECASE),
    # Fallback: "Server host key: <type> <fingerprint>"
    re.compile(r"Server host key:\s+(\S+)", re.IGNORECASE),
]
_PEER_SERVER_KEXINIT = re.compile(r"peer server KEXINIT proposal", re.IGNORECASE)


def _parse_host_key_algorithms(verbose_output: str) -> list[str]:
    """Extract host key algorithms from sftp -vvv stderr output.

    OpenSSH logs "host key algorithms:" for both the client's own KEXINIT
    proposal and the server's.  "peer server KEXINIT proposal" is logged by
    the OpenSSH client immediately before it parses the server's KEXINIT
    (regardless of remote server implementation), so we use it as an anchor
    and ignore any "host key algorithms:" match before it.

    Falls back to unambiguously server-side patterns if the anchor is absent.
    """
    seen_server_kexinit = False
    for line in verbose_output.splitlines():
        if not seen_server_kexinit:
            if _PEER_SERVER_KEXINIT.search(line):
                seen_server_kexinit = True
            continue

        m = _HOST_KEY_ALGO_PATTERN.search(line)
        if m:
            algos = [a.strip() for a in m.group(1).split(",") if a.strip()]
            if algos:
                return algos

    # Fall back to patterns that are unambiguously server-side
    for pattern in _SERVER_HOST_KEY_ALGO_PATTERNS:
        for line in verbose_output.splitlines():
            m = pattern.search(line)
            if m:
                algos = [a.strip() for a in m.group(1).split(",") if a.strip()]
                if algos:
                    return algos

    return []


def enumerate_algorithms(
    host: str, port: int, user: str, known_hosts_file: Path, host_key_algorithms: list[str]
) -> list[str]:
    section("Step 3: Enumerate server host key algorithms")

    cmd = [
        "sftp",
        "-vvv",
        "-o", "PreferredAuthentications=none",
        "-o", "StrictHostKeyChecking=yes",
        "-o", f"UserKnownHostsFile={known_hosts_file}",
        "-o", f"Port={port}",
        "-o", f"HostKeyAlgorithms={','.join(host_key_algorithms)}",
    ]
    cmd.append(f"{user}@{host}")

    print(f"$ {' '.join(cmd)}\n")

    try:
        result = run(cmd, timeout=20, input="\n")
    except subprocess.TimeoutExpired:
        fail("sftp timed out during algorithm enumeration.")

    output = result.stderr
    print(output.rstrip())

    if not output.strip():
        fail("sftp produced no output — connection may have been refused.")

    server_algos = _parse_host_key_algorithms(output)

    if server_algos:
        print(f"\n[OK] Host key algorithms offered by server: {', '.join(server_algos)}")
        compatible = set(host_key_algorithms) & set(server_algos)
        if compatible:
            print(f"[OK] Compatible algorithm(s): {', '.join(sorted(compatible))}")
        else:
            fail(
                f"No overlap between requested algorithms ({', '.join(host_key_algorithms)}) "
                f"and server's offered algorithms ({', '.join(server_algos)}).\n"
                "       Use --host-key-algorithms to specify a compatible algorithm."
            )
    else:
        print("\n[WARN] Could not parse host key algorithms from sftp output.")
        print("       Review the verbose output above manually.")

    return server_algos


# ---------------------------------------------------------------------------
# Step 4: authenticate with private key
# ---------------------------------------------------------------------------

def authenticate(
    host: str, port: int, user: str, key_path: str, known_hosts_file: Path, host_key_algorithms: list[str]
) -> None:
    section("Step 4: Authentication attempt with private key")

    cmd = [
        "sftp",
        "-vvv",
        "-o", "StrictHostKeyChecking=yes",
        "-o", f"UserKnownHostsFile={known_hosts_file}",
        "-o", "BatchMode=yes",
        "-o", f"Port={port}",
        "-o", f"HostKeyAlgorithms={','.join(host_key_algorithms)}",
        "-i", key_path,
    ]
    cmd.append(f"{user}@{host}")

    print(f"$ {' '.join(cmd)}\n")

    try:
        result = run(cmd, timeout=30, input="exit\n")
    except subprocess.TimeoutExpired:
        fail("sftp timed out during authentication.")

    output = result.stderr + result.stdout
    print(output.rstrip())

    if result.returncode == 0:
        print("\n[OK] Authentication succeeded.")
    else:
        if any(kw in output for kw in ("Permission denied", "Authentication failed")):
            fail("Authentication was rejected by the server. Check your key and username.")
        elif any(kw in output for kw in ("Connection refused", "Connection timed out", "No route")):
            fail("Connection failed during authentication step.")
        else:
            fail(f"sftp exited with code {result.returncode}. Review output above.")


# ---------------------------------------------------------------------------
# Paramiko mode: step 2 — host key fetch + verification
# ---------------------------------------------------------------------------

def paramiko_fetch_host_key(host: str, port: int, expected_key: str | None):
    """Fetch the server's host key via paramiko and verify it.

    Returns the paramiko PKey object for use in the auth step.
    """
    try:
        import paramiko
    except ImportError:
        fail("paramiko is not installed. Run: pip install paramiko")

    section(f"Step 2 (paramiko): Host key verification — {host}:{port}")

    transport = paramiko.Transport((host, port))
    try:
        transport.connect()
        server_key = transport.get_remote_server_key()
    except Exception as exc:
        fail(f"paramiko could not complete key exchange: {exc}")
    finally:
        transport.close()

    key_type = server_key.get_name()
    key_b64 = server_key.get_base64()
    print(f"[INFO] Server host key: {key_type} {key_b64}")

    scanned_keys = [f"{host} {key_type} {key_b64}"]
    if expected_key:
        _check_against_expected(scanned_keys, expected_key)
    else:
        _check_against_known_hosts(host, port, scanned_keys)

    return server_key


# ---------------------------------------------------------------------------
# Paramiko mode: step 3 — authenticate and open SFTP session
# ---------------------------------------------------------------------------

_PARAMIKO_KEY_CLASSES = None  # populated lazily after import


def _load_paramiko_key(key_path: str):
    """Try each paramiko key class in turn; return the first that succeeds."""
    import paramiko

    classes = [
        paramiko.Ed25519Key,
        paramiko.RSAKey,
        paramiko.ECDSAKey,
    ]
    for cls in classes:
        try:
            pkey = cls.from_private_key_file(key_path)
            print(f"[OK] Loaded private key as {cls.__name__}")
            return pkey
        except (paramiko.SSHException, ValueError):
            continue

    fail("Could not load private key — unsupported type or malformed key.")


def paramiko_connect(
    host: str,
    port: int,
    user: str,
    key_path: str,
    server_key,
    test_batch: bool = False,
) -> None:
    import paramiko

    section("Step 3 (paramiko): Authentication and SFTP session")

    pkey = _load_paramiko_key(key_path)

    client = paramiko.SSHClient()

    hostname_entry = f"[{host}]:{port}" if port != 22 else host
    client.get_host_keys().add(hostname_entry, server_key.get_name(), server_key)
    client.set_missing_host_key_policy(paramiko.RejectPolicy())

    print(f"[INFO] Connecting as {user}@{host}:{port} …")
    try:
        client.connect(host, port=port, username=user, pkey=pkey, timeout=30,
                       allow_agent=False, look_for_keys=False)
    except paramiko.AuthenticationException as exc:
        fail(f"Authentication rejected: {exc}")
    except paramiko.BadHostKeyException as exc:
        fail(f"Host key mismatch: {exc}")
    except paramiko.SSHException as exc:
        fail(f"SSH error during connect: {exc}")
    except OSError as exc:
        fail(f"Network error during connect: {exc}")

    print("[OK] SSH authentication succeeded.")

    try:
        sftp = client.open_sftp()
        cwd = sftp.normalize(".")
        print(f"[OK] SFTP session opened. Remote working directory: {cwd}")

        if test_batch:
            _paramiko_test_batch(sftp, cwd)

        sftp.close()
    except paramiko.SSHException as exc:
        fail(f"Could not open SFTP subsystem: {exc}")
    finally:
        client.close()


def _paramiko_test_batch(sftp, cwd: str) -> None:
    section("Step 4 (paramiko): SFTP batch operations test")

    test_dir = posixpath.join(cwd, f"sftp_debug_test_{os.getpid()}")
    print(f"[INFO] Test directory: {test_dir}")

    created = False
    try:
        sftp.mkdir(test_dir)
        created = True
        print(f"[OK] mkdir succeeded: {test_dir}")

        entries = sftp.listdir(cwd)
        dir_name = test_dir.split("/")[-1]
        if dir_name in entries:
            print(f"[OK] ls confirmed directory is visible in {cwd}")
        else:
            print(f"[WARN] mkdir succeeded but directory not found in ls output — entries: {entries}")
    except OSError as exc:
        fail(f"SFTP operation failed (mkdir {test_dir}): {exc}")
    finally:
        if created:
            try:
                sftp.rmdir(test_dir)
                print(f"[OK] rmdir succeeded: {test_dir}")
            except OSError as exc:
                print(f"[WARN] Could not remove test directory {test_dir}: {exc}", file=sys.stderr)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Step-by-step SFTP connection debugger.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python sftp_debug.py myuser example.com
  python sftp_debug.py myuser example.com 2222 --key-base64 "<base64-encoded-private-key>"
  python sftp_debug.py myuser example.com --key-base64 "<base64-encoded-private-key>" --expected-host-key "ssh-ed25519 AAAA..."
  python sftp_debug.py myuser example.com --key-base64 "<base64-encoded-private-key>" --host-key-algorithms ssh-rsa
  python sftp_debug.py myuser example.com --key-base64 "<base64-encoded-private-key>" --host-key-algorithms ssh-ed25519,ssh-rsa

  # paramiko mode (no sftp/ssh-keyscan binaries required)
  python sftp_debug.py myuser example.com --mode paramiko
  python sftp_debug.py myuser example.com 2222 --mode paramiko --key-base64 "<base64-encoded-private-key>"
  python sftp_debug.py myuser example.com --mode paramiko --key-base64 "<base64-encoded-private-key>" --expected-host-key "ssh-ed25519 AAAA..."
""",
    )
    parser.add_argument("user", help="SSH/SFTP username")
    parser.add_argument("host", help="Remote hostname or IP")
    parser.add_argument("port", type=int, nargs="?", default=22, help="SSH/SFTP port (default: 22)")
    parser.add_argument(
        "--key-base64",
        default=None,
        metavar="B64",
        help="Base64-encoded private key (optional; skips auth step if omitted).",
    )
    parser.add_argument(
        "--host-key-algorithms",
        default="ssh-ed25519,ssh-rsa",
        metavar="ALGOS",
        help="Comma-separated HostKeyAlgorithms to request (default: ssh-ed25519,ssh-rsa).",
    )
    parser.add_argument(
        "--expected-host-key",
        default=None,
        metavar="KEY",
        help=(
            "Expected host public key to verify against (e.g. 'ssh-ed25519 AAAA...'). "
            "If omitted, the script checks ~/.ssh/known_hosts instead."
        ),
    )
    parser.add_argument(
        "--mode",
        choices=["openssh", "paramiko"],
        default="openssh",
        help=(
            "Connection mode. 'openssh' (default) uses sftp/ssh-keyscan binaries. "
            "'paramiko' uses the paramiko Python library directly."
        ),
    )
    parser.add_argument(
        "--test-batch",
        action="store_true",
        default=False,
        help=(
            "After a successful paramiko connection, create a temporary directory, "
            "verify it appears in ls, then delete it. Only valid with --mode paramiko."
        ),
    )
    return parser.parse_args()


def _decode_key_to_tmpfile(key_base64: str) -> Path:
    """Decode a base64 private key and write it to a secure temp file."""
    try:
        key_bytes = base64.b64decode(key_base64)
    except Exception as exc:
        fail(f"Could not decode --key-base64 value: {exc}")
    tmp = tempfile.NamedTemporaryFile(
        prefix="sftp_debug_key_", delete=False, suffix=".pem"
    )
    tmp.write(key_bytes)
    tmp.close()
    os.chmod(tmp.name, 0o600)
    return Path(tmp.name)


def main() -> None:
    args = parse_args()

    check_reachability(args.host, args.port)

    if args.mode == "paramiko":
        server_key = paramiko_fetch_host_key(args.host, args.port, args.expected_host_key)
        tmp_key_file: Path | None = None
        try:
            if args.key_base64:
                tmp_key_file = _decode_key_to_tmpfile(args.key_base64)
                print(f"[INFO] Decoded private key written to temp file: {tmp_key_file}")
                paramiko_connect(args.host, args.port, args.user, str(tmp_key_file), server_key,
                                 test_ops=args.test_batch)
            else:
                print("\n[SKIP] No private key provided — skipping authentication step.")
        finally:
            if tmp_key_file and tmp_key_file.exists():
                tmp_key_file.unlink()
    else:
        host_key_algorithms = [a.strip() for a in args.host_key_algorithms.split(",") if a.strip()]
        known_hosts_file = keyscan(args.host, args.port, args.expected_host_key)
        enumerate_algorithms(args.host, args.port, args.user, known_hosts_file, host_key_algorithms)
        tmp_key_file = None
        try:
            if args.key_base64:
                tmp_key_file = _decode_key_to_tmpfile(args.key_base64)
                print(f"[INFO] Decoded private key written to temp file: {tmp_key_file}")
                authenticate(args.host, args.port, args.user, str(tmp_key_file), known_hosts_file, host_key_algorithms)
            else:
                print("\n[SKIP] No private key provided — skipping authentication step.")
        finally:
            if tmp_key_file and tmp_key_file.exists():
                tmp_key_file.unlink()

    print("\n[DONE] All steps passed.")


if __name__ == "__main__":
    main()