analyze_all_packages.py

import sys
import os
import argparse
import glob
import json
import csv
import analyze_package

def _discover_packages(traces_dir):
  """get all packages in trace directory; expects them to have been collected by pipeline.py

  filenames are all formatted as: {package_name}_{date}_{time}.pcap[ng]
  Returns: dict mapping package_name to list of files
  """
  packages = {}
  for filename in glob.glob(os.path.join(traces_dir, "*.pcap*")):
    basename = os.path.basename(filename)
    # strip the two trailing _date and _time segments
    package_name, _date, _time = basename.rsplit("_", 2)
    if package_name not in packages:
      packages[package_name] = []
    packages[package_name].append(filename)
  return packages

def main():
  parser = argparse.ArgumentParser(description="Analyze traces for all packages found in traces dir.")
  parser.add_argument("--dir", default="traces/", type=str)
  parser.add_argument("--json", default="analysis.jsonl", type=str)
  args = parser.parse_args()

  packages = _discover_packages(args.dir)
  if not packages:
    print(f"No pcap files found in {args.dir}", file=sys.stderr)
    sys.exit(1)

  writer = csv.writer(sys.stdout)
  writer.writerow(["package_name", "candidates_for_proprietary_crypto"])

  with open(args.json, "w") as f_out:
    for package_name, traces in sorted(packages.items()):
      try:
        results = analyze_package.analyze_all_files(traces)
      except Exception as e:
        writer.writerow([package_name, "ERROR"])
        continue

      total = sum(len(x) for x in results.values())
      json.dump({package_name: results}, f_out)
      f_out.write("\n")
      writer.writerow([package_name, str(total)])
      f_out.flush()
      sys.stdout.flush()


if __name__ == "__main__":
  main()