From 90cdef52327d3e8f110e769a5cb4d0e56e00f725 Mon Sep 17 00:00:00 2001 From: William Ballenthin Date: Thu, 13 Aug 2020 17:23:36 -0600 Subject: [PATCH] main: progress bar updates (+rules, and realize iterators) --- capa/main.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/capa/main.py b/capa/main.py index f5a50749..8a3dae68 100644 --- a/capa/main.py +++ b/capa/main.py @@ -106,7 +106,7 @@ def find_capabilities(ruleset, extractor, disable_progress=None): meta = {"feature_counts": {"file": 0, "functions": {},}} - for f in tqdm.tqdm(extractor.get_functions(), disable=disable_progress, unit=" functions"): + for f in tqdm.tqdm(list(extractor.get_functions()), disable=disable_progress, desc="matching", unit=" functions"): function_matches, bb_matches, feature_count = find_function_capabilities(ruleset, extractor, f) meta["feature_counts"]["functions"][f.__int__()] = feature_count logger.debug("analyzed function 0x%x and extracted %d features", f.__int__(), feature_count) @@ -315,7 +315,7 @@ def is_nursery_rule_path(path): return "nursery" in path -def get_rules(rule_path): +def get_rules(rule_path, disable_progress=False): if not os.path.exists(rule_path): raise IOError("rule path %s does not exist or cannot be accessed" % rule_path) @@ -343,7 +343,8 @@ def get_rules(rule_path): rule_paths.append(rule_path) rules = [] - for rule_path in rule_paths: + + for rule_path in tqdm.tqdm(list(rule_paths), disable=disable_progress, desc="loading ", unit=" rules"): try: rule = capa.rules.Rule.from_yaml_file(rule_path) except capa.rules.InvalidRule: @@ -526,7 +527,7 @@ def main(argv=None): logger.debug("using rules path: %s", rules_path) try: - rules = get_rules(rules_path) + rules = get_rules(rules_path, disable_progress=args.quiet) rules = capa.rules.RuleSet(rules) logger.debug("successfully loaded %s rules", len(rules)) if args.tag: