tor-browser

generate_builtins.py (6166B)

---

import argparse
import csv
import json
import struct
import subprocess
import sys
from datetime import datetime
from io import StringIO

import requests
from cryptography import x509
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization

# Format: CCADB Record Creation Date, SHA-256 Fingerprint, Subject Key Identifier, Authority Key Identifier, Root or Intermediate Certificate Record, X.509 Certificate PEM
REPORT_URL = "https://ccadb.my.salesforce-sites.com/ccadb/WebTrustListAsOf?ListDate={}"
DATE_ADDITION_COL = "CCADB Record Creation Date"
CERT_PEM_COL = "X.509 Certificate PEM"


class IdentifierAllocator:
   # Prefix should be a byte string.p
   def __init__(self, prefix):
       self.prefix = prefix
       self.position = 0

   def getIdentifier(self):
       result = self.prefix + struct.pack(">H", self.position)
       self.position += 1
       return result


def get_webtrust_certs(list_date):
   output = []
   url = REPORT_URL.format(list_date)
   response = requests.get(url, timeout=10)
   response.raise_for_status()
   csv_data = response.text
   with open(f"webtrust_certs-{list_date}.csv", "w") as certs_file:
       certs_file.write(csv_data)
   for r in csv.DictReader(StringIO(csv_data)):
       timestamp = datetime.strptime(r[DATE_ADDITION_COL], "%Y-%m-%dT%H:%M:%SZ")
       cert = x509.load_pem_x509_certificate(
           r[CERT_PEM_COL].encode("ascii"), default_backend()
       )
       certDer = cert.public_bytes(serialization.Encoding.DER)
       output.append((timestamp, certDer))
   print(f"Loaded {len(output)} certs from {url}")
   return output


def create_cert_dict(certs):
   certs.sort(key=lambda x: x[0])
   output = dict()
   idAlloc = IdentifierAllocator(b"\xff")
   for _, der in certs:
       idHex = idAlloc.getIdentifier()
       output[idHex] = der
   return output


def load_json_cache(jf):
   with open(jf) as f:
       j = json.load(f)
   assert j["data"]
   assert j["list_date"]
   assert j["creation_date"]
   j["data"] = {bytes.fromhex(k): bytes.fromhex(v) for k, v in j["data"].items()}
   return j


def cert_to_hash_rust_array(cert):
   digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
   digest.update(cert)

   # Finalize the hash and get the digest
   sha256_hash = digest.finalize()

   # Format the hash as a Rust array
   rust_array = ", ".join(f"0x{byte:02x}" for byte in sha256_hash)
   rust_output = f"[{rust_array}]"
   return rust_output


def make_rust_file_contents(certs, generation_date, list_date):
   output = ""
   output += f"""
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at https://mozilla.org/MPL/2.0/. */

use log;
use std::sync::OnceLock;
use thin_vec::ThinVec;

// Autogenerated via security/manager/ssl/abridged_certs/tools/generate_builtins.py
// Generation Date: {generation_date}
// Based on list version: {list_date}
"""
   output += """

// Public Interface

/// Given an Abridged Cert Identifier, lookup the hash of the corresponding certificate
pub fn id_to_hash(id: &[u8; 3]) -> Option<&ThinVec<u8>> {
   let index: usize = u16::from_be_bytes([id[1], id[2]]).into();
   log::trace!("Parsed identifier {:#02X?} as index {}", id, index);
   ABRIDGED_CERTS_BUILTINS_HASHES.get_or_init(init_hashes).get(index)
}

/// Get a list of hashes needed for this Abridged Certs scheme
pub fn get_needed_hashes() -> Option<&'static ThinVec<ThinVec<u8>>> {
   Some(ABRIDGED_CERTS_BUILTINS_HASHES.get_or_init(init_hashes))
}

// Private Implementation

/// This is currently built whenever it is accessed (similar a to lazy_static)
/// However, we may want explicit control in the future , e.g. if we want to delay
/// construction until after cert_storage has synced, or if we want to use a manifest
/// from remote settings.
static ABRIDGED_CERTS_BUILTINS_HASHES: OnceLock<ThinVec<ThinVec<u8>>> = OnceLock::new();

/// rustc / LLVM has a number of outstanding bugs in its code generation for large
/// functions. See the discussion in Bug 1969383 for why this format was selected.
"""
   output += f"const ABRIDGED_CERT_BYTES: [u8; {len(certs) * 32}] = ["
   for id, cert in certs.items():
       output += cert_to_hash_rust_array(cert).strip("[]") + ", "
   output += "];\n"

   output += """pub fn init_hashes() -> ThinVec<ThinVec<u8>>{"""
   output += f"""let mut m = ThinVec::with_capacity({len(certs)});"""
   output += """
           for entry in ABRIDGED_CERT_BYTES.chunks(32) {
               m.push(ThinVec::from(entry));
           }
           m.shrink_to_fit();
           m
       }
"""

   return output


if __name__ == "__main__":
   today = datetime.now().strftime("%Y-%m-%d")
   parser = argparse.ArgumentParser(
       description="Builds a map from identifiers to WebPKI Intermediate and Root Certificates"
   )
   parser.add_argument(
       "-d",
       "--date",
       help="Specify the date you want the list as-of (YYYY-MM-DD format)",
       type=str,
       default=today,
   )
   parser.add_argument(
       "-o",
       "--output",
       help="Specify the output file path",
       type=str,
       default="builtins.rs",
   )
   parser.add_argument(
       "-i",
       "--input",
       help="Specify a cached version of the list from JSON. Overrides date option.",
       type=str,
       default=None,
   )
   args = parser.parse_args()

   certs = None
   if args.input:
       j = load_json_cache(args.input)
       certs = j["data"]
       args.date = j["list_date"]
       today = j["creation_date"]
   else:
       certs = get_webtrust_certs(args.date)
       print(f"Fetched {len(certs)} certificates")
       certs = create_cert_dict(certs)

   with open(args.output, "w") as rust_file:
       rust_file.write(make_rust_file_contents(certs, today, args.date))

   subprocess.run(["rustfmt", args.output], capture_output=True, text=True, check=True)

   print(f"Generated file output to {args.output}")
   sys.exit(0)