import { lookup } from 'node:dns/promises';
import type { LookupAddress } from 'node:dns';
import net from 'node:net';

export class SsrfError extends Error {
  constructor(message: string) {
    super(message);
    this.name = 'SsrfError';
  }
}

function isPrivateIpv4(ip: string): boolean {
  const parts = ip.split('.').map(Number);
  if (parts.length !== 4 || parts.some((n) => Number.isNaN(n))) return true;
  const [a, b] = parts;
  if (a === 10) return true;
  if (a === 127) return true;
  if (a === 0) return true;
  if (a === 169 && b === 254) return true;
  if (a === 172 && b >= 16 && b <= 31) return true;
  if (a === 192 && b === 168) return true;
  if (a === 100 && b >= 64 && b <= 127) return true;
  if (a === 192 && b === 0) return true;
  if (a === 198 && (b === 18 || b === 19)) return true;
  if (a >= 224) return true;
  return false;
}

function isPrivateIpv6(ip: string): boolean {
  const lower = ip.toLowerCase();
  if (lower === '::1' || lower === '::') return true;
  if (lower.startsWith('fc') || lower.startsWith('fd')) return true;
  if (lower.startsWith('fe80')) return true;
  if (lower.startsWith('ff')) return true;
  // IPv4-mapped
  const v4 = lower.match(/^::ffff:(.+)$/);
  if (v4) return isPrivateIpv4(v4[1]);
  return false;
}

export function isPrivateIp(ip: string): boolean {
  if (net.isIP(ip) === 4) return isPrivateIpv4(ip);
  if (net.isIP(ip) === 6) return isPrivateIpv6(ip);
  return true;
}

export async function assertSafeUrl(rawUrl: string): Promise<void> {
  let url: URL;
  try {
    url = new URL(rawUrl);
  } catch {
    throw new SsrfError('invalid URL');
  }
  if (url.protocol !== 'http:' && url.protocol !== 'https:') {
    throw new SsrfError('protocol not allowed');
  }
  const host = url.hostname;
  // IP literal
  if (net.isIP(host)) {
    if (isPrivateIp(host)) throw new SsrfError('private address blocked');
    return;
  }
  let addresses: LookupAddress[];
  try {
    addresses = await lookup(host, { all: true });
  } catch {
    throw new SsrfError('dns resolution failed');
  }
  if (!addresses.length) throw new SsrfError('no addresses');
  for (const a of addresses) {
    if (isPrivateIp(a.address)) throw new SsrfError('private address blocked');
  }
}