#!/usr/bin/env python3
"""
Regenerate extension/shared/psl.js from the Public Suffix List.

    python scripts/update_psl.py

The extension has no build step — it is zipped as-is — so the list is vendored
as a plain classic script rather than pulled in via npm/bundler.

Why the extension needs this at all: autofill decides whether a stored
credential belongs to the page you are on. Plain suffix comparison treats
`evil.github.io` and `victim.github.io` as the same site, because `github.io`
looks like an ordinary domain. The PSL is the only way to know that it is a
public suffix and that those are different sites.

BOTH sections are included, deliberately:
  ICANN   — real TLDs (co.uk, com.au, ...)
  PRIVATE — github.io, vercel.app, herokuapp.com, ...

The PRIVATE section is the one that matters most here: those are the hosts where
an attacker can actually obtain a neighbouring subdomain.

Re-run when the list goes stale (it changes a few times a month). The generated
file records the upstream VERSION header so staleness is visible in a diff.
"""
import pathlib
import re
import sys
import urllib.request

PSL_URL = 'https://publicsuffix.org/list/public_suffix_list.dat'
OUT = pathlib.Path(__file__).resolve().parent.parent / 'extension' / 'shared' / 'psl.js'

HEADER = '''/**
 * extension/shared/psl.js — GENERATED FILE, DO NOT EDIT BY HAND.
 *
 * Regenerate with:  python scripts/update_psl.py
 *
 * Vendored Public Suffix List (https://publicsuffix.org/), used to decide
 * whether two hostnames belong to the same site before offering a stored
 * credential for autofill.
 *
 * Source list version: {version}
 * Rules: {n_rules} exact, {n_wild} wildcard, {n_exc} exception
 *
 * The list is MPL-2.0 licensed; see https://mozilla.org/MPL/2.0/.
 *
 * Exposes a single global, PkPsl, with:
 *   getRegistrableDomain(host) -> "example.co.uk" | null
 *   isSameSite(hostA, hostB)   -> boolean
 */
'''

BODY = r'''
const PkPsl = (() => {
  "use strict";

  // Split from single strings rather than array literals — same data, far less
  // punctuation, and the parse cost is a one-off at script load.
  const RULES = new Set(EXACT_BLOB.split("\n"));
  const WILDCARDS = new Set(WILD_BLOB ? WILD_BLOB.split("\n") : []);
  const EXCEPTIONS = new Set(EXC_BLOB ? EXC_BLOB.split("\n") : []);

  const IPV4_RE = /^\d{1,3}(\.\d{1,3}){3}$/;

  function _normalise(host) {
    if (!host) return null;
    let h = String(host).trim().toLowerCase();
    if (h.endsWith(".")) h = h.slice(0, -1);   // trailing root dot
    return h || null;
  }

  /**
   * Number of trailing labels that form the public suffix of `labels`.
   * Implements the matching rules from https://publicsuffix.org/list/:
   * exception rules win outright, otherwise the longest match wins, and an
   * unmatched host falls back to the implicit "*" rule.
   */
  function _publicSuffixLength(labels) {
    // Exception rules (!foo.bar) take priority over everything else.
    for (let i = 0; i < labels.length; i++) {
      if (EXCEPTIONS.has(labels.slice(i).join("."))) {
        return labels.length - i - 1;
      }
    }

    let best = 0;
    for (let i = 0; i < labels.length; i++) {
      const len = labels.length - i;
      if (len <= best) continue;
      if (RULES.has(labels.slice(i).join("."))) {
        best = len;
        continue;
      }
      // A wildcard rule "*.x.y" matches when labels[i] is any single label and
      // the remainder equals "x.y".
      if (i + 1 <= labels.length - 1 &&
          WILDCARDS.has(labels.slice(i + 1).join("."))) {
        best = len;
      }
    }

    // No rule matched: the implicit "*" rule makes the rightmost label the
    // public suffix (so "example.invalidtld" is still a registrable domain).
    return best || 1;
  }

  /**
   * The registrable domain ("example.co.uk") for a hostname, or null when the
   * host has none — an IP address, a single label like "localhost", or a host
   * that IS a public suffix ("github.io" itself).
   *
   * Callers must treat null as "no site identity": fall back to exact hostname
   * equality rather than assuming a match.
   */
  function getRegistrableDomain(host) {
    const h = _normalise(host);
    if (!h) return null;
    if (IPV4_RE.test(h) || h.includes(":")) return null;   // IPv4 / IPv6
    const labels = h.split(".");
    if (labels.length < 2) return null;                    // "localhost"

    const suffixLen = _publicSuffixLength(labels);
    if (labels.length <= suffixLen) return null;           // host is itself a suffix
    return labels.slice(labels.length - suffixLen - 1).join(".");
  }

  /**
   * True when two hostnames belong to the same registrable site.
   *
   * When either host has no registrable domain (IP, localhost, or a bare public
   * suffix) this falls back to exact hostname equality — never to a suffix
   * comparison, which is what allowed evil.github.io to match victim.github.io.
   */
  function isSameSite(a, b) {
    const ha = _normalise(a);
    const hb = _normalise(b);
    if (!ha || !hb) return false;
    if (ha === hb) return true;

    const da = getRegistrableDomain(ha);
    const db = getRegistrableDomain(hb);
    if (!da || !db) return false;
    return da === db;
  }

  return { getRegistrableDomain, isSameSite };
})();

// Content scripts and the popup load this as a classic script; the service
// worker imports it via importScripts. Export only where a module system exists.
if (typeof module !== "undefined" && module.exports) {
  module.exports = PkPsl;
}
'''


def main():
    print(f'Fetching {PSL_URL} ...')
    with urllib.request.urlopen(PSL_URL, timeout=60) as resp:
        text = resp.read().decode('utf-8')

    version = 'unknown'
    m = re.search(r'^// VERSION:\s*(.+)$', text, re.M)
    if m:
        version = m.group(1).strip()

    exact, wildcards, exceptions = [], [], []
    for line in text.splitlines():
        rule = line.strip()
        if not rule or rule.startswith('//'):
            continue
        if rule.startswith('!'):
            exceptions.append(rule[1:])
        elif rule.startswith('*.'):
            wildcards.append(rule[2:])
        elif '*' in rule:
            # No such rules exist today (wildcards are always leftmost). Skip
            # loudly rather than silently mis-parsing if that ever changes.
            print(f'  WARNING: skipping unsupported rule {rule!r}', file=sys.stderr)
        else:
            exact.append(rule)

    if len(exact) < 5000:
        sys.exit(f'FAIL: only {len(exact)} exact rules parsed — the list looks truncated')
    for expected in ('github.io', 'vercel.app', 'co.uk'):
        if expected not in exact:
            sys.exit(f'FAIL: expected rule {expected!r} missing — parse is wrong')

    def blob(name, values):
        return f'  const {name} = `' + '\n'.join(sorted(set(values))) + '`;\n'

    out = HEADER.format(version=version, n_rules=len(exact),
                        n_wild=len(wildcards), n_exc=len(exceptions))
    out += '\n// eslint-disable-next-line no-unused-vars\n'
    out += 'const _PSL_DATA = (() => {\n'
    out += blob('EXACT_BLOB', exact)
    out += blob('WILD_BLOB', wildcards)
    out += blob('EXC_BLOB', exceptions)
    out += '  return { EXACT_BLOB, WILD_BLOB, EXC_BLOB };\n})();\n'
    out += '\nconst { EXACT_BLOB, WILD_BLOB, EXC_BLOB } = _PSL_DATA;\n'
    out += BODY

    OUT.write_text(out, encoding='utf-8', newline='\n')
    size_kb = OUT.stat().st_size / 1024
    print(f'Wrote {OUT.relative_to(OUT.parent.parent.parent)} '
          f'({size_kb:.0f} KB) — version {version}')
    print(f'  {len(exact)} exact, {len(wildcards)} wildcard, {len(exceptions)} exception rules')


if __name__ == '__main__':
    main()