/**
 * Attack-vector tests for the /discover federation routes.
 *
 * These tests verify that the security fixes are working correctly.
 * Each describe block targets a vulnerability that was found during audit
 * and has now been patched.
 * 
 * @author Tocka
 * These tests were generated by AI (Claude Opus 4.6).
 * I did review it and made some changes but it's still mostly AI generated so take this with a grain of salt.
 * 
 * The AI did expose some actual vulnerabilities so it's not all bad.
 * After fixing the vulnerabilitie, the tests were updated to match the new behavior.
 * -----
 */
import { expect, test } from "@playwright/test"
import forge from "node-forge"
import http from "node:http"
import {
	clearTables,
	generateKeypair,
	getBlacklistedServer,
	getChallengesByServerUrl,
	seedChallenge,
	seedServer,
} from "./helpers/db"

const BASE = "http://localhost:3000"

function encryptPayload(payload: string, recipientPublicKeyPem: string) {
	const pub = forge.pki.publicKeyFromPem(recipientPublicKeyPem)
	return forge.util.encode64(
		pub.encrypt(forge.util.encodeUtf8(payload), "RSA-OAEP"),
	)
}

function fingerprintKey(pem: string): string {
	const md = forge.md.sha256.create()
	md.update(pem, "utf8")
	return md.digest().toHex()
}

function generate4096Keypair() {
	const kp = forge.pki.rsa.generateKeyPair(4096)
	return {
		publicKey: forge.pki.publicKeyToPem(kp.publicKey),
		privateKey: forge.pki.privateKeyToPem(kp.privateKey),
	}
}

function createTrapServer(fakePublicKey: string) {
	const hits: { method: string; url: string }[] = []
	const server = http.createServer((req, res) => {
		hits.push({ method: req.method!, url: req.url! })
		res.writeHead(200, { "Content-Type": "application/json" })
		res.end(JSON.stringify({ publicKey: fakePublicKey }))
	})

	return {
		hits,
		start: () => new Promise<number>((resolve) => {
			server.listen(0, "127.0.0.1", () => {
				resolve((server.address() as { port: number }).port)
			})
		}),
		stop: () => new Promise<void>((resolve) => { server.close(() => resolve()) }),
	}
}

test.beforeEach(async () => { await clearTables() })
test.afterEach(async () => { await clearTables() })

// ---------------------------------------------------------------------------
// 1. SSRF Protection — internal URLs are now blocked
// ---------------------------------------------------------------------------
test.describe("SSRF protection", () => {
	test("REGISTER rejects loopback URLs", async ({ request }) => {
		const { publicKey: fakePub } = generateKeypair()
		const trap = createTrapServer(fakePub)
		const port = await trap.start()

		try {
			const res = await request.post(`${BASE}/discover`, {
				data: {
					method: "REGISTER",
					url: `http://127.0.0.1:${port}`,
					publicKey: fakePub,
				},
			})

			expect(res.status()).toBe(400)
			const body = await res.json()
			expect(body.error).toMatch(/blocked/i)

			// The trap server should NOT have been hit
			expect(trap.hits.length).toBe(0)
		} finally {
			await trap.stop()
		}
	})

	test("REGISTER rejects RFC-1918 and link-local URLs", async ({ request }) => {
		const internalUrls = [
			"http://10.0.0.1:8080",
			"http://192.168.1.1:8080",
			"http://169.254.169.254",
		]

		for (const url of internalUrls) {
			const { publicKey } = generateKeypair()
			const res = await request.post(`${BASE}/discover`, {
				data: { method: "REGISTER", url, publicKey },
			})

			expect(res.status()).toBe(400)
			const body = await res.json()
			expect(body.error).toMatch(/blocked/i)
		}
	})

	test("DISCOVER rejects stored internal URLs", async ({ request }) => {
		const { publicKey: maliciousPub } = generateKeypair()
		await seedServer("http://127.0.0.1:9999", maliciousPub)

		// Build a valid signature using the fingerprint approach
		const signaturePayload = JSON.stringify({
			publicKeyFingerprint: fingerprintKey(maliciousPub),
			url: "http://127.0.0.1:9999",
		})
		const signature = encryptPayload(signaturePayload, process.env.FEDERATION_PUBLIC_KEY!)

		const res = await request.post(`${BASE}/discover`, {
			data: {
				method: "DISCOVER",
				publicKey: maliciousPub,
				signature,
			},
		})

		// Should be blocked rather than fetching the internal URL
		expect(res.status()).toBe(400)
		const body = await res.json()
		expect(body.error).toMatch(/blocked/i)
	})
})

// ---------------------------------------------------------------------------
// 2. Blacklist enforcement — blocks all federation routes
// ---------------------------------------------------------------------------
test.describe("Blacklist enforcement (fixed)", () => {
	async function blacklistServer(serverUrl: string, request: any) {
		// Seed a challenge with 1 attempt left
		await seedChallenge({
			serverUrl,
			attemptsLeft: 1,
			expiresAt: new Date(Date.now() + 1000 * 60 * 5),
		})

		// First attempt: mismatch, decrements to 0
		await request.post(`${BASE}/discover/rotate/confirm`, {
			data: {
				serverUrl,
				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
			},
		})

		// Second attempt: attemptsLeft=0, triggers blacklist
		await request.post(`${BASE}/discover/rotate/confirm`, {
			data: {
				serverUrl,
				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
			},
		})

		const bl = await getBlacklistedServer(serverUrl)
		expect(bl).toBeDefined()
	}

	test("blacklisted server is rejected by rotate/init", async ({ request }) => {
		const { publicKey: oldPub } = generateKeypair()
		const serverUrl = "https://blacklisted-server.example"
		await seedServer(serverUrl, oldPub)
		await blacklistServer(serverUrl, request as any)

		const { publicKey: newPub } = generate4096Keypair()
		const initRes = await request.post(`${BASE}/discover/rotate/init`, {
			data: { url: serverUrl, newPublicKey: newPub },
		})
		expect(initRes.status()).toBe(403)
		const body = await initRes.json()
		expect(body.error).toMatch(/blacklisted/i)
	})

	test("blacklisted server is rejected by rotate/confirm", async ({ request }) => {
		const serverUrl = "https://blacklisted-confirm.example"
		const { publicKey } = generateKeypair()
		await seedServer(serverUrl, publicKey)
		await blacklistServer(serverUrl, request as any)

		const confirmRes = await request.post(`${BASE}/discover/rotate/confirm`, {
			data: {
				serverUrl,
				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
			},
		})
		expect(confirmRes.status()).toBe(403)
		const body = await confirmRes.json()
		expect(body.error).toMatch(/blacklisted/i)
	})
})

// ---------------------------------------------------------------------------
// 3. Race condition fixed — transaction + FOR UPDATE
// ---------------------------------------------------------------------------
test.describe("Race condition fixed on rotate/confirm", () => {
	test("concurrent requests are serialised by the row lock", async () => {
		const serverUrl = "https://race-target.example"
		const { publicKey } = generateKeypair()
		await seedServer(serverUrl, publicKey)

		await seedChallenge({
			serverUrl,
			attemptsLeft: 1,
			expiresAt: new Date(Date.now() + 1000 * 60 * 5),
		})

		const payload = JSON.stringify({
			serverUrl,
			signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
			signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
		})

		const fire = () =>
			fetch(`${BASE}/discover/rotate/confirm`, {
				method: "POST",
				headers: { "Content-Type": "application/json" },
				body: payload,
			})

		const results = await Promise.all(Array.from({ length: 15 }, fire))
		const statuses = results.map((r) => r.status)

		const mismatch = statuses.filter((s) => s === 400).length
		const blacklisted = statuses.filter((s) => s === 403).length
		const notFound = statuses.filter((s) => s === 404).length

		// With the transaction lock, exactly 1 request should process the
		// mismatch (400), then the next sees attemptsLeft=0 and blacklists (403),
		// and the rest find no challenge (404) or hit the blacklist check.
		// The key point: no more than 1 mismatch (400) should get through.
		expect(mismatch).toBeLessThanOrEqual(1)
		expect(mismatch + blacklisted + notFound).toBe(statuses.length)
	})
})

// ---------------------------------------------------------------------------
// 4. Challenge deduplication — no flooding, no attempt reset
// ---------------------------------------------------------------------------
test.describe("Challenge deduplication (fixed)", () => {
	test("second init is rejected while a challenge is pending", async ({ request }) => {
		const serverUrl = "https://dedup-target.example"
		const { publicKey } = generateKeypair()
		await seedServer(serverUrl, publicKey)

		const { publicKey: newPub1 } = generate4096Keypair()
		const { publicKey: newPub2 } = generate4096Keypair()

		const res1 = await request.post(`${BASE}/discover/rotate/init`, {
			data: { url: serverUrl, newPublicKey: newPub1 },
		})
		expect(res1.status()).toBe(200)

		// Second init while the first is still active → 409
		const res2 = await request.post(`${BASE}/discover/rotate/init`, {
			data: { url: serverUrl, newPublicKey: newPub2 },
		})
		expect(res2.status()).toBe(409)
		const body = await res2.json()
		expect(body.error).toMatch(/already pending/i)

		// Only one challenge exists
		const challenges = await getChallengesByServerUrl(serverUrl)
		expect(challenges.length).toBe(1)
	})

	test("init succeeds after the previous challenge expires", async ({ request }) => {
		const serverUrl = "https://dedup-expire.example"
		const { publicKey } = generateKeypair()
		await seedServer(serverUrl, publicKey)

		// Seed an already-expired challenge directly
		await seedChallenge({
			serverUrl,
			expiresAt: new Date(Date.now() - 1000),
		})

		const { publicKey: newPub } = generate4096Keypair()
		const res = await request.post(`${BASE}/discover/rotate/init`, {
			data: { url: serverUrl, newPublicKey: newPub },
		})
		expect(res.status()).toBe(200)

		// Old challenge was replaced, only new one exists
		const challenges = await getChallengesByServerUrl(serverUrl)
		expect(challenges.length).toBe(1)
		expect(challenges[0].newPublicKey).toBe(newPub)
	})

	test("blacklisted server cannot reset attempts via new init", async ({ request }) => {
		const serverUrl = "https://reset-blocked.example"
		const { publicKey } = generateKeypair()
		await seedServer(serverUrl, publicKey)

		// Exhaust attempts → get blacklisted
		await seedChallenge({
			serverUrl,
			attemptsLeft: 1,
			expiresAt: new Date(Date.now() + 1000 * 60 * 5),
		})
		await request.post(`${BASE}/discover/rotate/confirm`, {
			data: {
				serverUrl,
				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
			},
		})
		await request.post(`${BASE}/discover/rotate/confirm`, {
			data: {
				serverUrl,
				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
			},
		})

		const bl = await getBlacklistedServer(serverUrl)
		expect(bl).toBeDefined()

		// Try init → blocked by blacklist check
		const { publicKey: freshPub } = generate4096Keypair()
		const initRes = await request.post(`${BASE}/discover/rotate/init`, {
			data: { url: serverUrl, newPublicKey: freshPub },
		})
		expect(initRes.status()).toBe(403)
	})
})

// ---------------------------------------------------------------------------
// 5. Signature validation — field values must match the request
// ---------------------------------------------------------------------------
test.describe("Signature validation (fixed)", () => {
	test("signature with mismatched publicKey fingerprint is rejected", async ({ request }) => {
		const { publicKey: peerPub } = generateKeypair()
		await seedServer("https://sig-test.example", peerPub)

		// Encrypt a signature where the fingerprint doesn't match
		const badSignature = encryptPayload(
			JSON.stringify({ publicKeyFingerprint: "wrong-fingerprint", url: "https://sig-test.example" }),
			process.env.FEDERATION_PUBLIC_KEY!,
		)

		const res = await request.post(`${BASE}/discover`, {
			data: {
				method: "DISCOVER",
				publicKey: peerPub,
				signature: badSignature,
			},
		})

		expect(res.status()).toBe(400)
	})

	test("signature with placeholder values is rejected", async ({ request }) => {
		const { publicKey: peerPub } = generateKeypair()
		await seedServer("https://sig-test2.example", peerPub)

		// The old bypass: { publicKey: "x", url: "y" }, now invalid
		const forgerySignature = encryptPayload(
			JSON.stringify({ publicKey: "x", url: "y" }),
			process.env.FEDERATION_PUBLIC_KEY!,
		)

		const res = await request.post(`${BASE}/discover`, {
			data: {
				method: "DISCOVER",
				publicKey: peerPub,
				signature: forgerySignature,
			},
		})

		expect(res.status()).toBe(400)
	})

	test("signature with correct fingerprint passes validation", async ({ request }) => {
		const { publicKey: peerPub } = generateKeypair()
		const trap = createTrapServer(peerPub)
		const port = await trap.start()
		const peerUrl = `http://127.0.0.1:${port}`

		try {
			await seedServer(peerUrl, peerPub)

			const validSignature = encryptPayload(
				JSON.stringify({
					publicKeyFingerprint: fingerprintKey(peerPub),
					url: peerUrl,
				}),
				process.env.FEDERATION_PUBLIC_KEY!,
			)

			const res = await request.post(`${BASE}/discover`, {
				data: {
					method: "DISCOVER",
					publicKey: peerPub,
					signature: validSignature,
				},
			})

			// The signature is valid, but the stored URL is internal → blocked by SSRF guard
			expect(res.status()).toBe(400)
			const body = await res.json()
			expect(body.error).toMatch(/blocked/i)
		} finally {
			await trap.stop()
		}
	})
})

// ---------------------------------------------------------------------------
// 6. Information disclosure: only url + isHealthy in peer list
// ---------------------------------------------------------------------------
test.describe("Information disclosure", () => {
	test("GET /discover only returns url and isHealthy for peers", async ({ request }) => {
		const { publicKey: peerPub1 } = generateKeypair()
		const { publicKey: peerPub2 } = generateKeypair()
		await seedServer("https://peer-one.example", peerPub1)
		await seedServer("https://peer-two.example", peerPub2)

		const res = await request.get(`${BASE}/discover`)
		expect(res.status()).toBe(200)

		const body = await res.json()
		expect(body.peers).toBeInstanceOf(Array)
		expect(body.peers.length).toBeGreaterThanOrEqual(2)

		for (const peer of body.peers) {
			expect(peer.url).toBeDefined()
			expect(peer.isHealthy).toBeDefined()
			// Internal fields must NOT be exposed
			expect(peer.id).toBeUndefined()
			expect(peer.createdAt).toBeUndefined()
			expect(peer.updatedAt).toBeUndefined()
			expect(peer.lastSeen).toBeUndefined()
			expect(peer.isHealthy).toBeUndefined()
		}
	})
})