diff --git a/polypyus/cli.py b/polypyus/cli.py index 2292448..98a8e98 100644 --- a/polypyus/cli.py +++ b/polypyus/cli.py @@ -21,6 +21,7 @@ ) from polypyus.graph import Graph from polypyus.importer import get_or_create_annotation, get_or_create_binary +from polypyus.exporter import export_csv_combined from polypyus.models import DB, Binary, Function, Match, Matcher from polypyus.tools import format_addr, format_data, format_percentage, serialize @@ -118,13 +119,15 @@ def format_binary_list( binary_list: Iterable[Binary], is_annotated: Optional[bool] = False ) -> Iterable[dict]: - keys = ["id", "name", "filepath"] + keys = ["id", "name", "filepath", "comment"] data = [] for binary in list(binaries): b_dict = binary.to_dict(only=keys) if is_annotated is True: b_dict["#annotations"] = binary.annotations.count() b_dict["#functions"] = binary.functions.count() + else: # this is a target binary + b_dict["#matches"] = binary.matches.count() data.append(b_dict) return data @@ -147,6 +150,18 @@ def format_binary_list( typer.echo(table) +@orm.db_session +@logger.catch +def _binary_ops(binary: str, comment: str, export_csv: str, remove: bool): + b = Binary.get(name=binary) + if comment is not None: + b.comment = comment + if export_csv is not None: + export_csv_combined(b, export_csv) + if remove is True: + Binary.delete(b) + + @app.command() @show_time def analyze( @@ -162,6 +177,12 @@ def analyze( False, help="List annotated binaries registered in project" ), list_targets: bool = typer.Option(False, help="List targets registered in project"), + binary: str = typer.Option("", help="Perform action on specific binary"), + comment: str = typer.Option(None, help="Add comment to binary"), + export_csv: str = typer.Option( + None, help="Export annotations/matches to specified csv file" + ), + remove: bool = typer.Option(False, help="remove binary from database"), ): """ Analyze targets with matchers generated from the given history (annotated binaries). @@ -179,6 +200,16 @@ def analyze( --min-size the minimum size in bytes a function needs to have to be considered for matcher creation. + --list-history/target lists the registered history or target binaries + + --binary specifies a specific binary to perform operations on + + --comment adds given comment to binary (requires --binary) + + --export-csv exports all matches/annotations for given binary to specified csv file + + --remove removes binary from database + """ if len(history) != len(annotation): @@ -193,6 +224,8 @@ def analyze( bind_db(project) if list_history is True or list_targets is True: _cli_list(list_history, list_targets) + elif binary != "": + _binary_ops(binary, comment, export_csv, remove) else: _analyze(history, annotation, target, parallelize, min_size, max_rel_fuzz) diff --git a/polypyus/exporter.py b/polypyus/exporter.py index 755d79c..e2e3e52 100644 --- a/polypyus/exporter.py +++ b/polypyus/exporter.py @@ -6,25 +6,44 @@ from polypyus.models import Binary from polypyus.tools import serialize +from typing import Iterable + csv.register_dialect("space_delimiter", delimiter=" ", quoting=csv.QUOTE_MINIMAL) +def export_csv_combined(binary: Binary, path: Path): + if binary.is_target is True: + export_matches_csv(binary, path) + else: + export_annotations_csv(binary, path) + + def export_matches_csv(binary: Binary, path: Path): logger.info(f"exporting matches for {binary.name} csv to {path}") stream = serialize(binary.matches, export=True) + _export_csv_internal(stream, path) + + +def export_annotations_csv(binary: Binary, path: Path): + logger.info(f"exporting annotations for {binary.name} csv to {path}") + stream = serialize(binary.functions, export=True) + _export_csv_internal(stream, path) + + +def _export_csv_internal(stream: Iterable[dict], path: Path): with open(path, "w") as csv_file: writer = csv.DictWriter( csv_file, fieldnames=CSV_KEYS_LONG, dialect="space_delimiter" ) writer.writeheader() - for match in stream: - match["addr"] = hex(match["addr"]) - match["name"] = match["name"].split(", ")[0] + for entry in stream: + entry["addr"] = hex(entry["addr"]) + entry["name"] = entry["name"].split(", ")[0] writer.writerow( { key: value - for key, value in {**match, "type": "FUNC"}.items() + for key, value in {**entry, "type": "FUNC"}.items() if key in CSV_KEYS_LONG } ) diff --git a/polypyus/importer.py b/polypyus/importer.py index 1dfcf25..ddcceec 100644 --- a/polypyus/importer.py +++ b/polypyus/importer.py @@ -106,7 +106,7 @@ def get_or_create_binary(path: Path, make_target=False) -> Binary: if not path.is_file(): raise FileNotFoundError binary = Binary(filepath=str(path), name=path.name) - elif make_target: + if make_target: binary.is_target = True binary.partition() return binary diff --git a/polypyus/models.py b/polypyus/models.py index 912b65d..9efb83c 100644 --- a/polypyus/models.py +++ b/polypyus/models.py @@ -67,7 +67,7 @@ class Binary(DB.Entity, Serializable): raw (bytes): binary blob, obtained by reading the binary functions (Set): Set of functions found in the binary matched_by (Set): Set of matches that matched something in this binary - + comment (str): A user defined comment for additional information """ name = Required(str) @@ -78,6 +78,7 @@ class Binary(DB.Entity, Serializable): functions = Set("Function") is_target = Required(bool, default=False) partitions = Optional(Json) + comment = Optional(str) @property def path_obj(self):