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

export type DnsResolver = (hostname: string) => Promise<string[]>;

export interface SafeUrlOptions {
	resolver?: DnsResolver;
	allowPrivateNetwork?: boolean;
}

export interface SafeUrlResult {
	url: URL;
	normalizedUrl: string;
	addresses: string[];
}

const BLOCKED_SUFFIXES = [".local", ".localhost", ".internal"];

export async function assertSafeHttpUrl(
	input: string | URL,
	options: SafeUrlOptions = {},
): Promise<SafeUrlResult> {
	const url = new URL(input.toString().trim());
	if (url.protocol !== "http:" && url.protocol !== "https:") {
		throw new Error(
			`Only http: and https: URLs are supported, got ${url.protocol}`,
		);
	}
	validateHostname(url, options.allowPrivateNetwork === true);
	if (options.allowPrivateNetwork === true) {
		return { url, normalizedUrl: url.toString(), addresses: [] };
	}
	const hostname = stripIpv6Brackets(url.hostname.toLowerCase());
	if (isIP(hostname) !== 0)
		return { url, normalizedUrl: url.toString(), addresses: [hostname] };
	const addresses = options.resolver
		? await options.resolver(hostname)
		: (await lookup(hostname, { all: true, verbatim: true })).map(
				(x) => x.address,
			);
	for (const address of addresses) {
		if (isPrivateOrReservedIp(address)) {
			throw new Error(
				`Hostname resolves to blocked private address: ${address}`,
			);
		}
	}
	return { url, normalizedUrl: url.toString(), addresses };
}

export async function resolveRedirectUrl(
	fromUrl: string,
	location: string,
	options: SafeUrlOptions = {},
): Promise<SafeUrlResult> {
	return assertSafeHttpUrl(new URL(location, fromUrl), options);
}

function validateHostname(url: URL, allowPrivate: boolean): void {
	const hostname = stripIpv6Brackets(url.hostname.toLowerCase());
	if (!hostname) throw new Error("URL must include a hostname");
	if (
		hostname === "localhost" ||
		BLOCKED_SUFFIXES.some((x) => hostname.endsWith(x))
	) {
		throw new Error(`Blocked private hostname: ${hostname}`);
	}
	if (!allowPrivate && isPrivateOrReservedIp(hostname)) {
		throw new Error(`Blocked private address: ${hostname}`);
	}
}

export function isPrivateOrReservedIp(address: string): boolean {
	const host = stripIpv6Brackets(address.toLowerCase());
	const family = isIP(host);
	if (family === 4) return isPrivateOrReservedIpv4(host);
	if (family === 6) return isPrivateOrReservedIpv6(host);
	return false;
}

function isPrivateOrReservedIpv4(address: string): boolean {
	const [a = 0, b = 0] = address.split(".").map((x) => Number.parseInt(x, 10));
	return (
		a === 0 ||
		a === 10 ||
		a === 127 ||
		(a === 169 && b === 254) ||
		(a === 172 && b >= 16 && b <= 31) ||
		(a === 192 && b === 168) ||
		(a === 100 && b >= 64 && b <= 127) ||
		a >= 224
	);
}

function isPrivateOrReservedIpv6(address: string): boolean {
	const embedded = address.match(
		/(?:::ffff:)?(\d{1,3}(?:\.\d{1,3}){3})$/u,
	)?.[1];
	if (embedded) return isPrivateOrReservedIpv4(embedded);
	const first = hextet(address, 0);
	const second = hextet(address, 1);
	const third = hextet(address, 2);
	const fourth = hextet(address, 3);
	return (
		address === "::" ||
		address === "::1" ||
		(first & 0xff00) === 0x0000 ||
		(first === 0x0100 && second === 0 && third === 0 && fourth === 0) ||
		(first & 0xfe00) === 0xfc00 ||
		(first & 0xffc0) === 0xfe80 ||
		(first & 0xff00) === 0xff00 ||
		first === 0x2002 ||
		(first === 0x2001 && second === 0x0db8)
	);
}

function hextet(address: string, index: number): number {
	return Number.parseInt(address.split(":")[index] || "0", 16) || 0;
}

function stripIpv6Brackets(hostname: string): string {
	return hostname.startsWith("[") && hostname.endsWith("]")
		? hostname.slice(1, -1)
		: hostname;
}