scripts: improve IDA scripts with direct annotation and Go/Rust support

scripts/idaplugin.py

@@ -39,6 +39,8 @@
 import floss.stackstrings
 import floss.tightstrings
 import floss.string_decoder
+from floss.language.go.extract import extract_go_strings
+from floss.language.rust.extract import extract_rust_strings
 from floss.results import AddressType, StackString, TightString, DecodedString
 
 logger = logging.getLogger("floss.idaplugin")
@@ -79,7 +81,9 @@ def append_comment(ea: int, s: str, repeatable: bool = False) -> None:
         idc.set_cmt(ea, cmt, False)
 
 
-def append_lvar_comment(fva: int, frame_offset: int, s: str, repeatable: bool = False) -> None:
+def append_lvar_comment(
+    fva: int, frame_offset: int, s: str, repeatable: bool = False
+) -> None:
     """
     add the given string as a (possibly repeatable) stack variable comment to the given function.
     does not add the comment if it already exists.
@@ -101,10 +105,15 @@ def append_lvar_comment(fva: int, frame_offset: int, s: str, repeatable: bool =
         idc.get_func_attr(fva, idc.FUNCATTR_FRSIZE) - frame_offset
     )  # alternative: idc.get_frame_lvar_size(fva) - frame_offset
     if not lvar_offset:
-        raise RuntimeError("failed to compute local variable offset: 0x%x 0x%x %s" % (fva, stack, s))
+        raise RuntimeError(
+            "failed to compute local variable offset: 0x%x 0x%x %s" % (fva, stack, s)
+        )
 
     if lvar_offset <= 0:
-        raise RuntimeError("failed to compute positive local variable offset: 0x%x 0x%x %s" % (fva, stack, s))
+        raise RuntimeError(
+            "failed to compute positive local variable offset: 0x%x 0x%x %s"
+            % (fva, stack, s)
+        )
 
     string = idc.get_member_cmt(stack, lvar_offset, repeatable)
     if not string:
@@ -115,7 +124,10 @@ def append_lvar_comment(fva: int, frame_offset: int, s: str, repeatable: bool =
         string = string + "\n" + s
 
     if not idc.set_member_cmt(stack, lvar_offset, string, repeatable):
-        raise RuntimeError("failed to set comment: 0x%08x 0x%08x 0x%08x: %s" % (fva, stack, lvar_offset, s))
+        raise RuntimeError(
+            "failed to set comment: 0x%08x 0x%08x 0x%08x: %s"
+            % (fva, stack, lvar_offset, s)
+        )
 
 
 def apply_decoded_strings(decoded_strings: List[DecodedString]) -> None:
@@ -124,15 +136,45 @@ def apply_decoded_strings(decoded_strings: List[DecodedString]) -> None:
             continue
 
         if ds.address_type == AddressType.GLOBAL:
-            logger.info("decoded string at global address 0x%x: %s", ds.address, ds.string)
+            logger.info(
+                "decoded string at global address 0x%x: %s", ds.address, ds.string
+            )
             append_comment(ds.address, ds.string)
         else:
-            logger.info("decoded string for function call at 0x%x: %s", ds.decoded_at, ds.string)
+            logger.info(
+                "decoded string for function call at 0x%x: %s", ds.decoded_at, ds.string
+            )
             append_comment(ds.decoded_at, ds.string)
 
 
+def apply_language_strings(fpath: Path, min_length: int = MIN_LENGTH) -> None:
+    """
+    Extract and apply language-specific strings for Go and Rust binaries.
+    """
+    language_extractors = {
+        "Go": extract_go_strings,
+        "Rust": extract_rust_strings,
+    }
+
+    for lang, extractor in language_extractors.items():
+        try:
+            strings = list(extractor(fpath, min_length))
+            if strings:
+                logger.info("extracted %d %s strings", len(strings), lang)
+                for s in strings:
+                    if s.string:
+                        append_comment(s.address, "FLOSS " + lang + ": " + s.string)
+            else:
+                logger.info("no %s strings found", lang)
+        except Exception as e:
+            logger.warning("failed to extract %s strings: %s", lang, str(e))
+
+
 def apply_stack_strings(
-    stack_strings: List[StackString], tight_strings: List[TightString], lvar_cmt: bool = True, cmt: bool = True
+    stack_strings: List[StackString],
+    tight_strings: List[TightString],
+    lvar_cmt: bool = True,
+    cmt: bool = True,
 ) -> None:
     """
     lvar_cmt: apply stack variable comment
@@ -144,7 +186,10 @@ def apply_stack_strings(
             continue
 
         logger.info(
-            "decoded stack/tight string in function 0x%x (pc: 0x%x): %s", s.function, s.program_counter, s.string
+            "decoded stack/tight string in function 0x%x (pc: 0x%x): %s",
+            s.function,
+            s.program_counter,
+            s.string,
         )
         if lvar_cmt:
             try:
@@ -185,19 +230,29 @@ def main(argv=None):
     time0 = time.time()
 
     logger.info("identifying decoding functions...")
-    decoding_function_features, library_functions = floss.identify.find_decoding_function_features(
-        vw, selected_functions, disable_progress=True
+    decoding_function_features, library_functions = (
+        floss.identify.find_decoding_function_features(
+            vw, selected_functions, disable_progress=True
+        )
     )
 
     logger.info("extracting stackstrings...")
-    selected_functions = floss.identify.get_functions_without_tightloops(decoding_function_features)
+    selected_functions = floss.identify.get_functions_without_tightloops(
+        decoding_function_features
+    )
     stack_strings = floss.stackstrings.extract_stackstrings(
-        vw, selected_functions, MIN_LENGTH, verbosity=floss.render.Verbosity.VERBOSE, disable_progress=True
+        vw,
+        selected_functions,
+        MIN_LENGTH,
+        verbosity=floss.render.Verbosity.VERBOSE,
+        disable_progress=True,
     )
     logger.info("decoded %d stack strings", len(stack_strings))
 
     logger.info("extracting tightstrings...")
-    tightloop_functions = floss.identify.get_functions_with_tightloops(decoding_function_features)
+    tightloop_functions = floss.identify.get_functions_with_tightloops(
+        decoding_function_features
+    )
     tight_strings = floss.tightstrings.extract_tightstrings(
         vw,
         tightloop_functions,
@@ -213,8 +268,12 @@ def main(argv=None):
 
     top_functions = floss.identify.get_top_functions(decoding_function_features, 20)
     fvas_to_emulate = floss.identify.get_function_fvas(top_functions)
-    fvas_tight_functions = floss.identify.get_tight_function_fvas(decoding_function_features)
-    fvas_to_emulate = floss.identify.append_unique(fvas_to_emulate, fvas_tight_functions)
+    fvas_tight_functions = floss.identify.get_tight_function_fvas(
+        decoding_function_features
+    )
+    fvas_to_emulate = floss.identify.append_unique(
+        fvas_to_emulate, fvas_tight_functions
+    )
     decoded_strings = floss.string_decoder.decode_strings(
         vw,
         fvas_to_emulate,
@@ -224,6 +283,8 @@ def main(argv=None):
     )
     logger.info("decoded %d strings", len(decoded_strings))
     apply_decoded_strings(decoded_strings)
+    logger.info("extracting language-specific strings (Go/Rust)...")
+    apply_language_strings(fpath)
 
     time1 = time.time()
     logger.debug("finished execution after %f seconds", (time1 - time0))

scripts/render-ida-import-script.py

@@ -40,6 +40,14 @@
 import argparse
 from pathlib import Path
 
+try:
+    import idc
+    import ida_kernwin
+
+    IDA_ENV = True
+except ImportError:
+    IDA_ENV = False
+
 from floss.results import AddressType, ResultDocument
 
 logger = logging.getLogger("floss.render-ida-import-script")
@@ -55,21 +63,30 @@ def render_ida_script(result_document: ResultDocument) -> str:
             b64 = base64.b64encode(ds.string.encode("utf-8")).decode("ascii")
             b64 = 'base64.b64decode("%s").decode("utf-8")' % b64
             if ds.address_type == AddressType.GLOBAL:
-                main_commands.append('print("FLOSS: string \\"%%s\\" at global VA 0x%x" %% (%s))' % (ds.address, b64))
-                main_commands.append('AppendComment(%d, "FLOSS: " + %s, True)' % (ds.address, b64))
+                main_commands.append(
+                    'print("FLOSS: string \\"%%s\\" at global VA 0x%x" %% (%s))'
+                    % (ds.address, b64)
+                )
+                main_commands.append(
+                    'AppendComment(%d, "FLOSS: " + %s, True)' % (ds.address, b64)
+                )
             else:
                 main_commands.append(
-                    'print("FLOSS: string \\"%%s\\" decoded at VA 0x%x" %% (%s))' % (ds.decoded_at, b64)
+                    'print("FLOSS: string \\"%%s\\" decoded at VA 0x%x" %% (%s))'
+                    % (ds.decoded_at, b64)
+                )
+                main_commands.append(
+                    'AppendComment(%d, "FLOSS: " + %s)' % (ds.decoded_at, b64)
                 )
-                main_commands.append('AppendComment(%d, "FLOSS: " + %s)' % (ds.decoded_at, b64))
     main_commands.append('print("Imported decoded strings from FLOSS")')
 
     for ss in result_document.strings.stack_strings:
         if ss.string != "":
             b64 = base64.b64encode(ss.string.encode("utf-8")).decode("ascii")
             b64 = 'base64.b64decode("%s").decode("utf-8")' % b64
             main_commands.append(
-                'AppendLvarComment(%d, %d, "FLOSS stackstring: " + %s, True)' % (ss.function, ss.frame_offset, b64)
+                'AppendLvarComment(%d, %d, "FLOSS stackstring: " + %s, True)'
+                % (ss.function, ss.frame_offset, b64)
             )
     main_commands.append('print("Imported stackstrings from FLOSS")')
 
@@ -78,7 +95,8 @@ def render_ida_script(result_document: ResultDocument) -> str:
             b64 = base64.b64encode(ts.string.encode("utf-8")).decode("ascii")
             b64 = 'base64.b64decode("%s").decode("utf-8")' % b64
             main_commands.append(
-                'AppendLvarComment(%d, %d, "FLOSS tightstring: " + %s, True)' % (ts.function, ts.frame_offset, b64)
+                'AppendLvarComment(%d, %d, "FLOSS tightstring: " + %s, True)'
+                % (ts.function, ts.frame_offset, b64)
             )
     main_commands.append('print("Imported tightstrings from FLOSS")')
 
@@ -132,15 +150,85 @@ def main():
     return script_content
 
 
+def run_in_ida():
+    """
+    Run directly as an IDAPython script inside IDA Pro.
+    Prompts user to select a FLOSS JSON result file and applies
+    annotations directly to the current IDB without generating
+    an intermediate script.
+    """
+    json_path = ida_kernwin.ask_file(0, "*.json", "Select FLOSS result JSON file")
+    if not json_path:
+        logger.warning("no file selected, exiting")
+        return
+
+    result_document = ResultDocument.parse_file(Path(json_path))
+
+    # apply decoded strings
+    for ds in result_document.strings.decoded_strings:
+        if not ds.string:
+            continue
+        if ds.address_type == AddressType.GLOBAL:
+            idc.set_cmt(ds.address, "FLOSS: " + ds.string, True)
+        else:
+            idc.set_cmt(ds.decoded_at, "FLOSS: " + ds.string, False)
+
+    # apply stack strings
+    for ss in result_document.strings.stack_strings:
+        if ss.string:
+            stack = idc.get_func_attr(ss.function, idc.FUNCATTR_FRAME)
+            if stack:
+                lvar_offset = (
+                    idc.get_func_attr(ss.function, idc.FUNCATTR_FRSIZE)
+                    - ss.frame_offset
+                )
+                if lvar_offset and lvar_offset > 0:
+                    idc.set_member_cmt(
+                        stack, lvar_offset, "FLOSS stackstring: " + ss.string, False
+                    )
+                    continue
+            # fallback to program counter if stack variable comment fails
+            idc.set_cmt(ss.program_counter, "FLOSS stackstring: " + ss.string, False)
+
+    # apply tight strings
+    for ts in result_document.strings.tight_strings:
+        if ts.string:
+            stack = idc.get_func_attr(ts.function, idc.FUNCATTR_FRAME)
+            if stack:
+                lvar_offset = (
+                    idc.get_func_attr(ts.function, idc.FUNCATTR_FRSIZE)
+                    - ts.frame_offset
+                )
+                if lvar_offset and lvar_offset > 0:
+                    idc.set_member_cmt(
+                        stack, lvar_offset, "FLOSS tightstring: " + ts.string, False
+                    )
+                    continue
+            # fallback to program counter if stack variable comment fails
+            idc.set_cmt(ts.program_counter, "FLOSS tightstring: " + ts.string, False)
+
+    ida_kernwin.refresh_idaview_anyway()
+    logger.info("successfully applied FLOSS annotations from %s", json_path)
+
+
 def main():
-    parser = argparse.ArgumentParser(description="Generate an IDA Python script to apply FLOSS results.")
-    parser.add_argument("/path/to/report.json", help="path to JSON document from `floss --json`")
+    parser = argparse.ArgumentParser(
+        description="Generate an IDA Python script to apply FLOSS results."
+    )
+    parser.add_argument(
+        "/path/to/report.json", help="path to JSON document from `floss --json`"
+    )
 
     logging_group = parser.add_argument_group("logging arguments")
 
-    logging_group.add_argument("-d", "--debug", action="store_true", help="enable debugging output on STDERR")
     logging_group.add_argument(
-        "-q", "--quiet", action="store_true", help="disable all status output except fatal errors"
+        "-d", "--debug", action="store_true", help="enable debugging output on STDERR"
+    )
+    logging_group.add_argument(
+        "-q",
+        "--quiet",
+        action="store_true",
+        help="disable all status output except fatal errors",
     )
 
     args = parser.parse_args()
@@ -163,4 +251,7 @@ def main():
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    if IDA_ENV:
+        run_in_ida()
+    else:
+        sys.exit(main())