Script to generate and write a basic TUF repo for profiling purposes

profile_tuf_bins.py

"""
Script to generate and write a basic TUF repo with keys, top-level roles, a
delegated targets role, and BIN_N_COUNT delegated targets roles, for profiling
purposes.

<Usage>

pip install securesystemslib[crypto,pynacl] tuf
curl https://gist.githubusercontent.com/lukpueh/724bd1d7b477f201a9f199b037d85747/raw/profile_tuf_bins.py -o profile_tuf_bins.py
python -m cProfile -o stats profile_tuf_bins.py
python
>>> import pstats, os, glob
>>> from pstats import SortKey
>>> newest_stats = os.path.join(sorted(glob.glob("tuf_stats_*")).pop(), "stats")
>>> p = pstats.Stats(newest_stats)
>>> p.sort_stats(SortKey.CUMULATIVE).print_stats("tuf|securesystemslib|copy", 20)

"""
import os
import time
from tuf import repository_tool

BIN_N_COUNT = 16384

PREFIX_LEN = len(f"{BIN_N_COUNT - 1:x}")
PREFIX_COUNT = 16 ** PREFIX_LEN
BIN_SIZE = PREFIX_COUNT // BIN_N_COUNT

TUF_REPO = "test-repo"
TOPLEVEL_ROLES = ["root", "snapshot", "targets", "timestamp"]
BINS_ROLE = "bins"
BIN_N_ROLE = "bin-n"

KEYS = {}

# Create and change into timestamped experiment directory
_dir = f"tuf_stats__{int(time.time())}__bins_{BIN_N_COUNT}"
os.mkdir(_dir)
os.chdir(_dir)

def main():
  # Create keys pairs for all roles
  for role in TOPLEVEL_ROLES + [BINS_ROLE] + [BIN_N_ROLE]:
    key_name = repository_tool.generate_and_write_ed25519_keypair(password="pw")
    private = repository_tool.import_ed25519_privatekey_from_file(key_name, password="pw")
    public = repository_tool.import_ed25519_publickey_from_file(key_name + ".pub")
    KEYS[role] = {"private": private, "public": public}

  # Generate TUF repo with top level roles
  repository = repository_tool.create_new_repository(TUF_REPO)
  for role in TOPLEVEL_ROLES:
    role_obj = getattr(repository, role)
    role_obj.threshold = 1
    role_obj.add_verification_key(KEYS[role]["public"])
    role_obj.load_signing_key(KEYS[role]["private"])

  # Delegate from "top-level targets" to "bins" targets role
  repository.targets.delegate(BINS_ROLE, [KEYS[BINS_ROLE]["public"]], [])
  repository.targets(BINS_ROLE).load_signing_key(KEYS[BINS_ROLE]["private"])

  # Delegate from "bins" to "bin-n" targets roles
  repository.targets(BINS_ROLE).delegate_hashed_bins(
      [], [KEYS[BIN_N_ROLE]["public"]], BIN_N_COUNT)

  # Mark all "bin-n" roles as dirty and write everything
  dirty_roles = []
  for idx in range(0, PREFIX_COUNT, BIN_SIZE):
    role = f"{idx:0{PREFIX_LEN}x}-{idx+BIN_SIZE-1:0{PREFIX_LEN}x}"
    dirty_roles.append(role)

  repository.mark_dirty(dirty_roles)
  repository.writeall(consistent_snapshot=True)

if __name__ == '__main__':
  main()