Allowing users to add tags to strings

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "tests/data"]
 	path = tests/data
-	url = ../flare-floss-testfiles.git
+	url = https://github.com/mandiant/flare-floss-testfiles.git

floss/qs/db/expert.py

@@ -82,8 +82,17 @@ def from_file(cls, path: pathlib.Path) -> "ExpertStringDatabase":
         )
 
 
-DEFAULT_PATHS = (pathlib.Path(floss.qs.db.__file__).parent / "data" / "expert" / "capa.jsonl",)
+DEFAULT_PATHS = (
+    pathlib.Path(floss.qs.db.__file__).parent / "data" / "expert" / "capa.jsonl",
+    pathlib.Path(floss.qs.db.__file__).parent / "data" / "expert" / "user.jsonl",
+)
 
+def create_user_db():
+    user_json = pathlib.Path(floss.qs.db.__file__).parent / "data" / "expert" / "user.jsonl"
+    if not user_json.exists():
+        user_json.parent.mkdir(parents=True, exist_ok=True)
+        user_json.write_text("")
 
 def get_default_databases() -> Sequence[ExpertStringDatabase]:
+    create_user_db()
     return [ExpertStringDatabase.from_file(path) for path in DEFAULT_PATHS]

floss/qs/main.py

@@ -8,6 +8,7 @@
 import hashlib
 import logging
 import pathlib
+import msgspec
 import argparse
 import datetime
 import functools
@@ -39,6 +40,17 @@
 
 
 QS_VERSION = "0.1.0"
+KNOWN_TAGS = {
+    "#code",
+    "#code-junk",
+    "#common",
+    "#duplicate",
+    "#reloc",
+    "#winapi",
+    "#decoded",
+    "#capa"
+}
+USER_DB_PATH = pathlib.Path(floss.qs.db.__file__).parent / "data" / "expert" / "user.jsonl"
 
 
 @contextlib.contextmanager
@@ -1317,6 +1329,54 @@ def render_string_line(console: Console, tag_rules: TagRules, string: ResultStri
 
         console.print(footer)
 
+def create_user_db():
+    if not USER_DB_PATH.exists():
+        USER_DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+        USER_DB_PATH.write_text("")
+
+def add_to_user_db(path, note, author, reference):
+    with open(path, 'r', encoding='utf-8') as f:
+        data = json.loads(f.read())
+        strings = collect_strings(data["layout"])
+        create_user_db()
+        new_entries = []
+        for s in strings:
+            unknown_tags = s.get("unknown_tags", [])
+            if not unknown_tags:
+                continue
+            for tag in unknown_tags:
+                new_string = {
+                    "type": "string",
+                    "value": s["string"],
+                    "tag": tag,
+                    "action": "highlight",
+                    "note": note,
+                    "description": "",
+                    "authors": [author] if author else [],
+                    "references": [r.strip() for r in reference.split(',')] if reference else []
+                }
+                new_entries.append(msgspec.json.encode(new_string).decode('utf-8'))
+
+        if new_entries:
+            with open(USER_DB_PATH, 'a', encoding='utf-8') as user_db:
+                user_db.write('\n'.join(new_entries) + '\n')
+
+def collect_strings(node, results = None):
+    if results is None:
+        results = []
+    if "strings" in node and node["strings"]:
+        for s in node["strings"]:
+            tags = s.get("tags", [])
+            unknownTags = [t for t in tags if t not in KNOWN_TAGS]
+            if unknownTags:
+                results.append({
+                    "string": s["string"],
+                    "unknown_tags": unknownTags
+                })
+    if "children" in node:
+        for child in node["children"]:
+            collect_strings(child, results)
+    return results
 
 def main():
     # set environment variable NO_COLOR=1 to disable color output.
@@ -1328,7 +1388,7 @@ def main():
         version=f"%(prog)s {QS_VERSION}",
         help="show program's version number and exit",
     )
-    parser.add_argument("path", help="file or path to analyze")
+    parser.add_argument("path", nargs="?", help="file or path to analyze")
     parser.add_argument(
         "-n",
         "--minimum-length",
@@ -1339,6 +1399,10 @@ def main():
     )
     parser.add_argument("-j", "--json", action="store_true", help="emit JSON instead of text")
     parser.add_argument("-l", "--load", action="store_true", help="load from existing FLOSS QUANTUMSTRAND results document")
+    parser.add_argument("--json-out", help="path to write layout to as JSON")
+    parser.add_argument("--json-in", help="path to read layout from as JSON")
+
+    parser.add_argument("--expand", "-e", nargs="?", const=True, help="add strings to database")
 
     logging_group = parser.add_argument_group("logging arguments")
     logging_group.add_argument("-d", "--debug", action="store_true", help="enable debugging output on STDERR")
@@ -1348,6 +1412,7 @@ def main():
         action="store_true",
         help="disable all status output except fatal errors",
     )
+
     args = parser.parse_args()
 
     floss.main.set_log_config(args.debug, args.quiet)
@@ -1367,15 +1432,54 @@ def main():
         sys.stdout.reconfigure(encoding="utf-8")
     colorama.just_fix_windows_console()
 
-    path = pathlib.Path(args.path)
-    if not path.exists():
-        logging.error("%s does not exist", path)
-        return 1
+    # Check if path is required based on the operation
+    if not args.load and not args.expand and not args.path:
+        parser.error("path argument is required when not using --load or --expand")
+
+    if args.path:
+        path = pathlib.Path(args.path)
+        if not path.exists():
+            logging.error("%s does not exist", path)
+            return 1
 
     if args.load:
-        with path.open("r") as f:
+        if args.path:
+            load_path = pathlib.Path(args.path)
+        else:
+            # If no path provided with --load, we need to get the JSON file path
+            if not args.json_in:
+                parser.error("--load requires either a path argument or --json-in option")
+            load_path = pathlib.Path(args.json_in)
+
+        if not load_path.exists():
+            logging.error("%s does not exist", load_path)
+            return 1
+
+        with load_path.open("r") as f:
             results = ResultDocument.model_validate_json(f.read())
+    elif args.expand:
+        if args.expand is True:
+            if not args.path:
+                parser.error("--expand without a value requires a path argument")
+            expand_path = pathlib.Path(args.path)
+        else:
+            expand_path = pathlib.Path(args.expand)
+
+        if not expand_path.exists():
+            logging.error("%s does not exist", expand_path)
+            return 1
+
+        note = input("A note for these strings: ")
+        author = input("Author: ")
+        reference = input("Reference: ")
+        add_to_user_db(str(expand_path), note, author, reference)
+        return 0
     else:
+        # Normal analysis mode - path is required
+        if not args.path:
+            parser.error("path argument is required for analysis")
+
+        path = pathlib.Path(args.path)
         with path.open("rb") as f:
             # because we store all the strings in memory
             # in order to tag and reason about them
@@ -1420,6 +1524,12 @@ def main():
         )
         results = ResultDocument.from_qs(meta, layout)
 
+    # Output handling - works for both load and analysis modes
+    if args.json_out:
+        with pathlib.Path(args.json_out).open("w") as f:
+            f.write(results.model_dump_json(indent=2))
+        logger.info("Wrote layout to %s", args.json_out)
+
     if args.json:
         print(results.model_dump_json(indent=0))
     else:
@@ -1431,9 +1541,10 @@ def main():
             "#reloc": "hide",
             # lib strings are muted (default)
         }
+
         # hide (remove) strings according to the above rules
         hide_strings_by_rules(results.layout, tag_rules)
-
+        
         console = Console()
         render_strings(console, results.layout, tag_rules)