/**
 * Shared SSRF (Server-Side Request Forgery) primitives.
 *
 * Originally lived in src/mcp/ssrf-strict.ts; extracted here so it can be
 * reused by SSH preflight (which connects to a raw host:port, not an HTTPS URL).
 *
 * src/mcp/ssrf-strict.ts is now a thin re-export — all logic lives here.
 *
 * Plan: docs/superpowers/plans/2026-05-12-ssh-tool-integration.md (Phase 2).
 */
import { promises as dns } from 'node:dns';
import * as net from 'node:net';
import type { LookupAddress } from 'node:dns';
import { Agent, fetch as undiciFetch } from 'undici';

export type LookupFn = (hostname: string) => Promise<Array<{ address: string; family: 4 | 6 }>>;

export type SsrfResult =
  | { ok: true; pinnedIp: string; family: 4 | 6 }
  | { ok: false; reason: string };

const defaultLookup: LookupFn = async (hostname) => {
  const addrs = await dns.lookup(hostname, { all: true });
  return addrs.map((a) => ({ address: a.address, family: a.family as 4 | 6 }));
};

function ipv4ToInt(ip: string): number | null {
  const parts = ip.split('.');
  if (parts.length !== 4) return null;
  let n = 0;
  for (const p of parts) {
    const v = Number(p);
    if (!Number.isInteger(v) || v < 0 || v > 255) return null;
    n = (n << 8) | v;
  }
  return n >>> 0;
}

function ipv4InCidr(ip: string, cidr: string): boolean {
  const [base, bitsStr] = cidr.split('/');
  const bits = Number(bitsStr);
  const ipN = ipv4ToInt(ip);
  const baseN = ipv4ToInt(base);
  if (ipN === null || baseN === null) return false;
  const mask = bits === 0 ? 0 : (~0 << (32 - bits)) >>> 0;
  return (ipN & mask) === (baseN & mask);
}

function normalizeIpv6(ip: string): string {
  return ip.toLowerCase();
}

function ipv6HasPrefix(ip: string, prefix: string): boolean {
  const p = prefix.toLowerCase();
  const addr = normalizeIpv6(ip);
  if (!net.isIPv6(addr)) return false;
  return addr.startsWith(p);
}

export function isPrivateOrForbidden(ip: string, family: 4 | 6): boolean {
  if (family === 4) {
    if (!ipv4ToInt(ip) && ip !== '0.0.0.0') return true; // malformed
    const cidrs = [
      '0.0.0.0/8',        // "this network"
      '10.0.0.0/8',       // RFC1918
      '100.64.0.0/10',    // CGNAT
      '127.0.0.0/8',      // loopback
      '169.254.0.0/16',   // link-local + IMDSv4
      '172.16.0.0/12',    // RFC1918
      '192.0.0.0/24',     // IETF protocol
      '192.168.0.0/16',   // RFC1918
      '198.18.0.0/15',    // benchmarking
      '224.0.0.0/4',      // multicast
      '240.0.0.0/4',      // reserved
      '255.255.255.255/32',
    ];
    return cidrs.some((c) => ipv4InCidr(ip, c));
  }

  // IPv6
  const addr = normalizeIpv6(ip);
  if (addr === '::' || addr === '::1') return true;
  if (ipv6HasPrefix(addr, '::ffff:')) return true;       // IPv4-mapped
  if (ipv6HasPrefix(addr, '64:ff9b:')) return true;      // NAT64
  if (ipv6HasPrefix(addr, 'fd00:ec2:')) return true;     // AWS IMDS IPv6
  if (/^f[cd]/.test(addr)) return true;                  // fc00::/7 ULA
  if (/^fe[89ab]/.test(addr)) return true;               // fe80::/10 link-local
  if (/^ff/.test(addr)) return true;                     // multicast
  return false;
}

export async function checkSSRFStrict(
  urlStr: string,
  opts: { lookup?: LookupFn } = {},
): Promise<SsrfResult> {
  let url: URL;
  try {
    url = new URL(urlStr);
  } catch {
    return { ok: false, reason: 'Invalid URL' };
  }
  if (url.protocol !== 'https:') {
    return { ok: false, reason: 'HTTPS required' };
  }

  const lookup = opts.lookup ?? defaultLookup;
  let addrs: Array<{ address: string; family: 4 | 6 }>;
  try {
    addrs = await lookup(url.hostname);
  } catch (err) {
    return { ok: false, reason: `DNS lookup failed: ${(err as Error).message}` };
  }
  if (addrs.length === 0) return { ok: false, reason: 'DNS resolution returned no addresses' };

  for (const a of addrs) {
    if (isPrivateOrForbidden(a.address, a.family)) {
      return { ok: false, reason: `Forbidden IP in DNS response: ${a.address}` };
    }
  }
  return { ok: true, pinnedIp: addrs[0].address, family: addrs[0].family };
}

export interface PinnedFetchOptions extends RequestInit {
  pinnedIp: string;
  family: 4 | 6;
}

export async function pinnedFetch(
  urlStr: string,
  opts: PinnedFetchOptions,
): Promise<Response> {
  const { pinnedIp, family, ...rest } = opts;
  // Node's net.connect calls lookup with options.all=true and expects an array of
  // LookupAddress; older callers use the (address, family) form. Support both.
  type LookupCb = (
    err: NodeJS.ErrnoException | null,
    address: string | LookupAddress[],
    family?: number,
  ) => void;
  const lookup = (
    _hostname: string,
    options: { all?: boolean },
    cb: LookupCb,
  ): void => {
    if (options && options.all === true) {
      cb(null, [{ address: pinnedIp, family }]);
    } else {
      cb(null, pinnedIp, family);
    }
  };
  const agent = new Agent({
    connect: {
      lookup: lookup as never,
    },
  });
  let res: Awaited<ReturnType<typeof undiciFetch>>;
  try {
    // undici fetch accepts `dispatcher`; cast because lib.dom RequestInit does not declare it.
    res = await undiciFetch(urlStr, { ...rest, dispatcher: agent } as never);
  } catch (err) {
    // Connection/request failed — nothing in flight, tear the agent down now.
    agent.destroy().catch(() => {});
    throw err;
  }
  // Graceful close: `Agent.close()` resolves only after in-flight requests
  // complete, i.e. after the response body has been fully consumed (or
  // cancelled / GC-finalized). Callers that stream the body — including
  // long-lived SSE transports (MCP) — are therefore not cut off, while the
  // per-call Agent no longer leaks its sockets once the body is done.
  agent.close().catch(() => {});
  // Convert undici Response into the web Response shape the callers already use.
  return res as unknown as Response;
}

// ── SSH-side primitives ─────────────────────────────────────────────────
//
// SSH connects to a raw host:port, not an HTTPS URL. The two helpers below
// give SSH the same DNS-pinning property: resolve once with policy enforcement,
// then connect to the resolved literal IP so name and target can't drift apart.

export interface ResolveCheckArgs {
  host: string;
  /** When true, private/loopback/link-local IPs pass policy. */
  allowPrivate: boolean;
  /** Test seam. */
  lookup?: LookupFn;
}

export type ResolveCheckResult =
  | { ok: true; ip: string; family: 4 | 6 }
  | { ok: false; reason: string };

/**
 * Resolve a hostname (or use a literal IP) and check policy.
 *
 * If the input is already a literal IPv4/IPv6, no DNS query is made; the
 * literal is checked directly. Otherwise DNS resolves all addresses and
 * every returned address must pass policy (`allowPrivate` toggle).
 *
 * The returned `ip` is the first resolved (or literal) address. SSH callers
 * must pass this same IP to `pinnedConnect` to defeat DNS rebinding.
 */
export async function resolveAndCheck(args: ResolveCheckArgs): Promise<ResolveCheckResult> {
  const { host, allowPrivate } = args;
  if (!host || typeof host !== 'string') {
    return { ok: false, reason: 'Invalid host' };
  }

  // Literal IP shortcut — no DNS query.
  const literalFamily = net.isIP(host);
  if (literalFamily === 4 || literalFamily === 6) {
    const fam = literalFamily as 4 | 6;
    if (!allowPrivate && isPrivateOrForbidden(host, fam)) {
      return { ok: false, reason: `Private/forbidden IP not allowed: ${host}` };
    }
    return { ok: true, ip: host, family: fam };
  }

  const lookup = args.lookup ?? defaultLookup;
  let addrs: Array<{ address: string; family: 4 | 6 }>;
  try {
    addrs = await lookup(host);
  } catch (err) {
    return { ok: false, reason: `DNS lookup failed: ${(err as Error).message}` };
  }
  if (addrs.length === 0) return { ok: false, reason: 'DNS resolution returned no addresses' };

  if (!allowPrivate) {
    for (const a of addrs) {
      if (isPrivateOrForbidden(a.address, a.family)) {
        return { ok: false, reason: `Forbidden IP in DNS response: ${a.address}` };
      }
    }
  }
  return { ok: true, ip: addrs[0].address, family: addrs[0].family };
}

export interface PinnedConnectArgs {
  /** Literal IP (already checked by resolveAndCheck). DNS rebinding defense. */
  ip: string;
  port: number;
  /** Connect deadline. Rejects with `Error('connect_timeout')` on expiry. */
  timeoutMs: number;
}

/**
 * Open a raw TCP socket to a pinned literal IP.
 *
 * The socket is configured with the connect timeout but the timer is cleared
 * once `connect` succeeds; the caller (ssh2) manages further I/O timeouts.
 * On timeout or error the socket is destroyed.
 */
export function pinnedConnect(args: PinnedConnectArgs): Promise<net.Socket> {
  return new Promise((resolve, reject) => {
    const family = net.isIP(args.ip);
    if (family !== 4 && family !== 6) {
      reject(new Error(`pinnedConnect: ip is not a literal address: ${args.ip}`));
      return;
    }
    const socket = new net.Socket();
    let done = false;
    const timer = setTimeout(() => {
      if (done) return;
      done = true;
      socket.destroy();
      reject(new Error('connect_timeout'));
    }, args.timeoutMs);
    socket.once('error', (err) => {
      if (done) return;
      done = true;
      clearTimeout(timer);
      socket.destroy();
      reject(err);
    });
    socket.once('connect', () => {
      if (done) return;
      done = true;
      clearTimeout(timer);
      resolve(socket);
    });
    socket.connect({ host: args.ip, port: args.port, family: family as 4 | 6 });
  });
}
