import { lookup } from 'node:dns/promises'
import { isIP } from 'node:net'

/**
 * Refuses to fetch anything that lives on this machine or this network.
 *
 * `fb_upload_image` takes a URL and the agent chooses it. On the production
 * server that ran inside a container; here it runs on the user's laptop, where
 * `127.0.0.1:18789` is their own gateway, `127.0.0.1:*` is whatever else they
 * have running, and `192.168.*` is their home network. Fetching such a URL and
 * uploading the bytes to Meta as "ad creative" is a working exfiltration
 * primitive, and the agent does not have to be malicious to be talked into it —
 * a prompt injection in any page or file it read is enough.
 *
 * So every hop is resolved and checked before it is opened, and redirects are
 * followed by hand rather than by `fetch`: a public URL that 302s to
 * `http://127.0.0.1` would otherwise walk straight through a check done only on
 * the first address.
 *
 * DNS rebinding is the residual gap — the name could resolve differently
 * between this check and the connection. Closing it properly needs a socket
 * bound to the address that was verified, which Node's `fetch` does not expose.
 * The window is small and the payoff for an attacker is one image fetch, so
 * this is documented rather than papered over.
 */

const MAX_REDIRECTS = 3

/** Ranges that are this machine, this network, or otherwise not the public internet. */
function isPrivateV4(address) {
  const parts = address.split('.').map(Number)
  if (parts.length !== 4 || parts.some((part) => Number.isNaN(part))) return true
  const [a, b] = parts
  if (a === 0) return true // "this network"
  if (a === 10) return true
  if (a === 127) return true // loopback
  if (a === 100 && b >= 64 && b <= 127) return true // carrier-grade NAT
  if (a === 169 && b === 254) return true // link-local, and cloud metadata
  if (a === 172 && b >= 16 && b <= 31) return true
  if (a === 192 && b === 168) return true
  if (a === 192 && b === 0) return true // IETF protocol assignments
  if (a === 198 && (b === 18 || b === 19)) return true // benchmarking
  if (a >= 224) return true // multicast and reserved
  return false
}

function isPrivateAddress(address) {
  const normalised = address.toLowerCase()
  if (isIP(normalised) === 4) return isPrivateV4(normalised)

  // IPv4 written as IPv6 — the v4 rules still apply, and this comes in two
  // spellings. `new URL()` rewrites `::ffff:127.0.0.1` to its hex form
  // `::ffff:7f00:1`, so checking only the dotted spelling let loopback through;
  // the test for it failed on exactly that, which is why both are handled.
  const mapped = normalised.match(/^::ffff:(.+)$/)
  if (mapped) {
    const tail = mapped[1]
    if (isIP(tail) === 4) return isPrivateV4(tail)
    const groups = tail.split(':')
    if (groups.length === 2) {
      const high = Number.parseInt(groups[0], 16)
      const low = Number.parseInt(groups[1], 16)
      if (Number.isFinite(high) && Number.isFinite(low)) {
        return isPrivateV4(
          `${(high >> 8) & 0xff}.${high & 0xff}.${(low >> 8) & 0xff}.${low & 0xff}`,
        )
      }
    }
    // An `::ffff:` address whose tail cannot be read is not one to trust.
    return true
  }

  if (normalised === '::' || normalised === '::1') return true
  if (/^f[cd]/.test(normalised)) return true // unique local
  if (/^fe[89ab]/.test(normalised)) return true // link local
  return false
}

/** Throws unless every address the host resolves to is on the public internet. */
export async function assertPublicUrl(raw) {
  let url
  try {
    url = new URL(raw)
  } catch {
    throw new Error(`Not a URL: ${raw}`)
  }
  if (url.protocol !== 'https:' && url.protocol !== 'http:') {
    throw new Error(`Only http and https URLs can be fetched, not ${url.protocol}`)
  }

  const host = url.hostname.replace(/^\[|\]$/g, '')
  const addresses = isIP(host) ? [{ address: host }] : await lookup(host, { all: true })
  if (addresses.length === 0) throw new Error(`Could not resolve ${host}`)

  for (const { address } of addresses) {
    if (isPrivateAddress(address)) {
      throw new Error(
        `Refusing to fetch ${host} — it resolves to ${address}, which is on this machine or this network.`,
      )
    }
  }
  return url
}

/**
 * `fetch`, with every redirect target checked the same way as the first URL.
 *
 * The response is returned unread so the caller decides what to do with the
 * body — this only guarantees where it came from.
 */
export async function fetchPublic(raw, init = {}) {
  let target = raw
  for (let hop = 0; hop <= MAX_REDIRECTS; hop++) {
    await assertPublicUrl(target)
    const response = await fetch(target, { ...init, redirect: 'manual' })
    if (response.status < 300 || response.status >= 400) return response

    const location = response.headers.get('location')
    if (!location) return response
    target = new URL(location, target).toString()
  }
  throw new Error(`Too many redirects fetching ${raw}`)
}
