Support downloading packages for other architectures

This is really handy for working with the armhf and arm64 systems we
build code for.

Signed-off-by: Brian Silverman <bsilver16384@gmail.com>
Change-Id: I966df8da77d6a985c4e3f5b1ad8b2b637bed3f98
diff --git a/debian/download_packages.py b/debian/download_packages.py
index e3bb845..e8caebd 100755
--- a/debian/download_packages.py
+++ b/debian/download_packages.py
@@ -2,23 +2,54 @@
 
 import sys
 import os
+import os.path
 import re
 import subprocess
 import tempfile
+import urllib.request
 import argparse
 import hashlib
 
-def get_deps(package):
+def initialize_apt(apt_dir, apt_args, args):
+  os.mkdir(os.path.join(apt_dir, 'etc'))
+  os.mkdir(os.path.join(apt_dir, 'etc', 'apt'))
+  os.mkdir(os.path.join(apt_dir, 'etc', 'apt', 'trusted.gpg.d'))
+  os.mkdir(os.path.join(apt_dir, 'etc', 'apt', 'preferences.d'))
+  os.mkdir(os.path.join(apt_dir, 'var'))
+  os.mkdir(os.path.join(apt_dir, 'var', 'lib'))
+  os.mkdir(os.path.join(apt_dir, 'var', 'lib', 'dpkg'))
+  with open(os.path.join(apt_dir, 'var', 'lib', 'dpkg', 'status'), 'w'):
+    pass
+  with open(os.path.join(apt_dir, 'etc', 'apt', 'sources.list'), 'w') as f:
+    f.write("""
+deb http://deb.debian.org/debian/ {release} main contrib non-free
+deb-src http://deb.debian.org/debian/ {release} main contrib non-free
+
+deb https://security.debian.org/debian-security {release}-security main contrib non-free
+deb-src https://security.debian.org/debian-security {release}-security main contrib non-free
+
+deb http://deb.debian.org/debian/ {release}-updates main contrib non-free
+deb-src http://deb.debian.org/debian/ {release}-updates main contrib non-free
+
+deb http://deb.debian.org/debian {release}-backports main contrib non-free
+deb-src http://deb.debian.org/debian {release}-backports main contrib non-free
+""".format(release=args.release))
+  for key in args.apt_key:
+    basename = os.path.basename(key)
+    urllib.request.urlretrieve(key, os.path.join(apt_dir, 'etc', 'apt', 'trusted.gpg.d', basename))
+  subprocess.check_call(["apt-get"] + apt_args + ["update"])
+
+def get_deps(apt_args, package):
   env = dict(os.environ)
   del env['LD_LIBRARY_PATH']
-  out = subprocess.check_output(["apt-rdepends", package], env=env)
+  out = subprocess.check_output(["apt-rdepends"] + apt_args + [package], env=env)
   deps = out.splitlines()
   return set([dep for dep in deps if not dep.startswith(b" ")])
 
-def get_all_deps(packages):
+def get_all_deps(apt_args, packages):
   deps = set()
   for package in packages or ():
-    deps.update(get_deps(package))
+    deps.update(get_deps(apt_args, package))
   return deps
 
 def map_virtual_packages(packages):
@@ -63,15 +94,15 @@
       continue
     yield package
 
-def download_deps(packages, excludes, force_includes):
-  deps = get_all_deps(packages)
-  exclude_deps = get_all_deps(excludes)
+def download_deps(apt_args, packages, excludes, force_includes):
+  deps = get_all_deps(apt_args, packages)
+  exclude_deps = get_all_deps(apt_args, excludes)
   deps -= exclude_deps
-  force_include_deps = get_all_deps(force_includes)
+  force_include_deps = get_all_deps(apt_args, force_includes)
   deps |= force_include_deps
   env = dict(os.environ)
   del env['LD_LIBRARY_PATH']
-  subprocess.check_call([b"apt-get", b"download"] + list(map_virtual_packages(deps)), env=env)
+  subprocess.check_call([b"apt-get"] + [a.encode('utf-8') for a in apt_args] + [b"download"] + list(map_virtual_packages(deps)), env=env)
 
 def fixup_files():
   # Gotta remove those pesky epoch numbers in the file names. Bazel doesn't
@@ -114,15 +145,34 @@
   parser = argparse.ArgumentParser()
   parser.add_argument("--exclude", "-e", type=str, action="append", help="A package to exclude from the list")
   parser.add_argument("--force-include", type=str, action="append", help="Force include this and its dependencies. Even if listed in excludes.")
+  parser.add_argument("--arch", type=str, default="amd64", help="Architecture to download files for.")
+  parser.add_argument("--apt-dir", type=str, help=" ".join([
+    "File to generate and store apt files in.",
+    "Helpful for saving time when downloading multiple groups of packages.",
+    "Some flags will be ignored in favor of the values used to create this folder, so be careful.",
+    ]))
+  parser.add_argument("--release", type=str, default="bullseye", help="Debian release to use.")
+  parser.add_argument("--apt-key", type=str, action="append", default=[
+    "https://ftp-master.debian.org/keys/archive-key-11.asc",
+    "https://ftp-master.debian.org/keys/archive-key-11-security.asc",
+  ], help="URL of an additional apt archive key to trust.")
   parser.add_argument("package", nargs="+", help="The packages to download.")
   args = parser.parse_args(argv[1:])
+  if args.apt_dir:
+    apt_dir = args.apt_dir
+  else:
+    apt_dir = tempfile.mkdtemp()
+  apt_args = ["-o", "Dir=" + apt_dir, "-o", "APT::Architecture=" + args.arch]
+  if not args.apt_dir:
+    print("Creating apt files in %s" % apt_dir)
+    initialize_apt(apt_dir, apt_args, args)
   folder = tempfile.mkdtemp()
   os.chdir(folder)
   excludes = args.exclude or []
   # Exclude common packages that don't make sense to include in everything all
   # the time.
   excludes += _ALWAYS_EXCLUDE
-  download_deps(args.package, excludes, args.force_include)
+  download_deps(apt_args, args.package, excludes, args.force_include)
   fixup_files()
   print_file_list()
   print("Your packages are all in %s" % folder)