cros_extract_deps: generate CPE only when user requests it

Generating CPEs is the slowest part (> 6 minutes) of the script. If
user specifies the 'deps' format, don't generate CPEs.

BUG=None
TEST=`cros_extract_deps --format=cpe --board=peppy virtual/target-os`
TEST=`cros_extract_deps --format=deps --board=peppy virtual/target-os`

Change-Id: I362d0c043c02e75ed39adbeba8018402b40809a2
Reviewed-on: https://chromium-review.googlesource.com/225176
Reviewed-by: Yu-Ju Hong <yjhong@chromium.org>
Commit-Queue: Yu-Ju Hong <yjhong@chromium.org>
Tested-by: Yu-Ju Hong <yjhong@chromium.org>

1 files changed, 50 insertions, 37 deletions

diff --git a/scripts/cros_extract_deps.py b/scripts/cros_extract_deps.py
index dab2c7f82..3e7d69b17 100644
--- a/scripts/cros_extract_deps.py
+++ b/scripts/cros_extract_deps.py
@@ -15,7 +15,7 @@ from chromite.lib import commandline
 from chromite.lib import cros_build_lib
 
 
-def FlattenDepTree(deptree, pkgtable=None, parentcpv=None):
+def FlattenDepTree(deptree, pkgtable=None, parentcpv=None, get_cpe=False):
   """Simplify dependency json.
 
 Turn something like this (the parallel_emerge DepsTree format):
@@ -53,28 +53,41 @@ Turn something like this (the parallel_emerge DepsTree format):
     "action": "merge"
   }
 }
+
+  Args:
+    deptree: The dependency tree.
+    pkgtable: The package table to update. If None, create a new one.
+    parentcpv: The parent CPV.
+    get_cpe: If set True, include CPE in the flattened dependency tree.
+
+  Returns:
+    A flattened dependency tree.
   """
   if pkgtable is None:
     pkgtable = {}
-  for cpv, record in deptree.items():
+  for cpv, record in deptree.iteritems():
     if cpv not in pkgtable:
       cat, nam, ver, rev = portage.versions.catpkgsplit(cpv)
-      pkgtable[cpv] = {"deps": [],
-                       "rev_deps": [],
-                       "name": nam,
-                       "category": cat,
-                       "version": "%s-%s" % (ver, rev),
-                       "full_name": cpv,
-                       "cpes": GetCPEFromCPV(cat, nam, ver),
-                       "action": record["action"]}
+      pkgtable[cpv] = {'deps': [],
+                       'rev_deps': [],
+                       'name': nam,
+                       'category': cat,
+                       'version': '%s-%s' % (ver, rev),
+                       'full_name': cpv,
+                       'cpes': [],
+                       'action': record['action']}
+      if get_cpe:
+        pkgtable[cpv]['cpes'].extend(GetCPEFromCPV(cat, nam, ver))
+
     # If we have a parent, that is a rev_dep for the current package.
     if parentcpv:
-      pkgtable[cpv]["rev_deps"].append(parentcpv)
+      pkgtable[cpv]['rev_deps'].append(parentcpv)
     # If current package has any deps, record those.
-    for childcpv in record["deps"]:
-      pkgtable[cpv]["deps"].append(childcpv)
+    for childcpv in record['deps']:
+      pkgtable[cpv]['deps'].append(childcpv)
     # Visit the subtree recursively as well.
-    FlattenDepTree(record["deps"], pkgtable=pkgtable, parentcpv=cpv)
+    FlattenDepTree(record['deps'], pkgtable=pkgtable, parentcpv=cpv,
+                   get_cpe=get_cpe)
   return pkgtable
 
 
@@ -90,7 +103,7 @@ def GetCPEFromCPV(category, package, version):
     A list of CPE Name strings, e.g.
     ["cpe:/a:curl:curl:7.30.0", "cpe:/a:curl:libcurl:7.30.0"]
   """
-  equery_cmd = ["equery", "m", "-U", "%s/%s" % (category, package)]
+  equery_cmd = ['equery', 'm', '-U', '%s/%s' % (category, package)]
   lines = cros_build_lib.RunCommand(equery_cmd, error_code_ok=True,
                                     print_cmd=False,
                                     redirect_stdout=True).output.splitlines()
@@ -98,9 +111,9 @@ def GetCPEFromCPV(category, package, version):
   # and extract the cpe URI.
   cpes = []
   for line in lines:
-    if "ID: cpe" not in line:
+    if 'ID: cpe' not in line:
       continue
-    cpes.append("%s:%s" % (line.split()[1], version.replace("_", "")))
+    cpes.append('%s:%s' % (line.split()[1], version.replace('_', '')))
   # Note that we're assuming we can combine the root of the CPE, taken
   # from metadata.xml, and tack on the version number as used by
   # Portage, and come up with a legitimate CPE. This works so long as
@@ -128,42 +141,42 @@ def GetCPEFromCPV(category, package, version):
 
 def ExtractCPEList(deps_list):
   cpe_dump = []
-  for cpv, record in deps_list.items():
-    if record["cpes"]:
-      name = "%s/%s" % (record["category"], record["name"])
-      cpe_dump.append({"ComponentName": name,
-                       "Repository": "cros",
-                       "Targets": sorted(record["cpes"])})
+  for cpv, record in deps_list.iteritems():
+    if record['cpes']:
+      name = '%s/%s' % (record['category'], record['name'])
+      cpe_dump.append({'ComponentName': name,
+                       'Repository': 'cros',
+                       'Targets': sorted(record['cpes'])})
     else:
-      cros_build_lib.Warning("No CPE entry for %s", cpv)
-  return sorted(cpe_dump, key=lambda k: k["ComponentName"])
+      cros_build_lib.Warning('No CPE entry for %s', cpv)
+  return sorted(cpe_dump, key=lambda k: k['ComponentName'])
 
 
 def main(argv):
   parser = commandline.ArgumentParser(description="""
 This extracts the dependency tree for the specified package, and outputs it
 to stdout, in a serialized JSON format.""")
-  parser.add_argument("--board", required=True,
-                      help="The board to use when computing deps.")
-  parser.add_argument("--format", default="deps",
-                      choices=["deps", "cpe"],
-                      help="Output either traditional deps or CPE-only JSON")
+  parser.add_argument('--board', required=True,
+                      help='The board to use when computing deps.')
+  parser.add_argument('--format', default='deps',
+                      choices=['deps', 'cpe'],
+                      help='Output either traditional deps or CPE-only JSON')
   # Even though this is really just a pass-through to DepGraphGenerator,
   # handling it as a known arg here allows us to specify a default more
   # elegantly than testing for its presence in the unknown_args later.
-  parser.add_argument("--root-deps", default="rdeps",
-                      help="Which deps to report (defaults to rdeps)")
+  parser.add_argument('--root-deps', default='rdeps',
+                      help='Which deps to report (defaults to rdeps)')
   known_args, unknown_args = parser.parse_known_args(argv)
 
-  lib_argv = ["--board=%s" % known_args.board,
-              "--root-deps=%s" % known_args.root_deps,
-              "--quiet", "--pretend", "--emptytree"]
+  lib_argv = ['--board=%s' % known_args.board,
+              '--root-deps=%s' % known_args.root_deps,
+              '--quiet', '--pretend', '--emptytree']
   lib_argv.extend(unknown_args)
 
   deps = DepGraphGenerator()
   deps.Initialize(lib_argv)
   deps_tree, _deps_info = deps.GenDependencyTree()
-  deps_list = FlattenDepTree(deps_tree)
-  if known_args.format == "cpe":
+  deps_list = FlattenDepTree(deps_tree, get_cpe=(known_args.format == 'cpe'))
+  if known_args.format == 'cpe':
     deps_list = ExtractCPEList(deps_list)
   print(json.dumps(deps_list, sort_keys=True, indent=2))