Script to find go dependency paths to bad library versions

deps.py

from collections import defaultdict
import sys
import semver
import pprint

def parse_version(input):
  return semver.VersionInfo.parse(input.lstrip('v'))

def is_bad_version(input, badlib, badver):
  try:
    lib, verstr = input.split('@')
  except ValueError:
    return False
  ver = parse_version(verstr)
  return lib == badlib and ver <= badver

def dfs(G, node, path):
  paths = []
  if not G.get(node):
    path = [node] + path
    return [path]
  for dst in G.get(node):
    paths.extend(dfs(G, dst, [node] + path))
  return paths

def mkgraphs(edges):
  adj_lists = defaultdict(lambda: set())
  inv_adj_lists = defaultdict(lambda: set())
  for edge in edges:
    for i in range(len(edge)-1):
      src = edge[i]
      dst = edge[i+1]

      adj_lists[src].add(dst)
      inv_adj_lists[dst].add(src)
  return adj_lists, inv_adj_lists

def print_trie(T, node, indent=""):
  print(f"{indent}{node}")
  if not indent:
    indent = "=> "
  for child in sorted(T.get(node) or []):
    print_trie(T, child, "  " + indent)

def main(args):
  if len(args) != 2:
    print("Usage: go mod graph | python3 deps.py github.com/bad/library@v1.2.3")
    sys.exit(1)

  badlibstr = args[1]
  badlib, badverstr = badlibstr.split('@')
  badver = parse_version(badverstr)

  input_str = ''.join([segment for segment in sys.stdin])

  edge_strings = [line for line in input_str.splitlines()]
  edge_pairs = [edge.split() for edge in edge_strings]
  # Ensure that each edge has 2 elements (src and dst)
  assert(all([len(e) == 2 for e in edge_pairs]))

  root = edge_pairs[0][0]
  print(f"Root package: {root}")

  bad_versions = set()

  _, inv_adj_lists = mkgraphs(edge_pairs)
  for node in inv_adj_lists.keys():
    if is_bad_version(node, badlib, badver):
      bad_versions.add(node)

  print(f"Bad libs found:")
  pprint.pprint(bad_versions)

  paths = []
  for bad_version in bad_versions:
    paths.extend(dfs(inv_adj_lists, bad_version, []))

  trie, _ = mkgraphs(paths)

  print(f"{len(paths)} paths to bad libs found:")
  print_trie(trie, root)

if __name__ == '__main__':
  main(sys.argv)