Add flush_to_disk() call in script when user downloads resource in GUI (

#277)

Address #274

frontend/src/ResourceTreeToolbar.svelte

@@ -129,6 +129,7 @@
             a.click();
 
             URL.revokeObjectURL(blobUrl);
+            await $selectedResource.add_flush_to_disk_to_script(a.download);
           }
         },
       },

frontend/src/ofrak/remote_resource.js

@@ -486,6 +486,21 @@ export class RemoteResource extends Resource {
       script.set(await r.json());
     });
   }
+
+  async add_flush_to_disk_to_script(output_file_name) {
+    await fetch(`${this.uri}/add_flush_to_disk_to_script`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify(output_file_name),
+    }).then(async (r) => {
+      if (!r.ok) {
+        throw Error(JSON.stringify(await r.json(), undefined, 2));
+      }
+      await this.update_script();
+    });
+  }
 }
 
 export function remote_models_to_resources(remote_models, resources) {

frontend/src/ofrak/resource.js

@@ -386,6 +386,10 @@ export class Resource {
     throw new NotImplementedError("update_script");
   }
 
+  async add_flush_to_disk_to_script() {
+    throw new NotImplementedError("add_flush_to_disk_to_script");
+  }
+
   /***
    * Return the string representation of a comment, which includes its range as prefix if it
    * has one.

ofrak_core/CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [Unreleased](https://github.com/redballoonsecurity/ofrak/tree/master)
 ### Added
+- Add call to flush a resource to disk in the script whenever a user downloads a resource from the GUI. ([#277](https://github.com/redballoonsecurity/ofrak/pull/277))
 - Generate dynamic, runnable script based on GUI actions and display the script in the GUI. ([#265](https://github.com/redballoonsecurity/ofrak/pull/265))
 - Add `-f`/`--file` option to `ofrak gui` command to pre-load some files into OFRAK before opening the GUI, so they can be explored right away ([#266](https://github.com/redballoonsecurity/ofrak/pull/266))
 - Add `-i`/`--import` option to the CLI to import and discover additional OFRAK Python packages when starting OFRAK. [#269](https://github.com/redballoonsecurity/ofrak/pull/269)

ofrak_core/ofrak/gui/server.py

@@ -154,6 +154,9 @@ def __init__(
                 web.post("/{resource_id}/delete_comment", self.delete_comment),
                 web.post("/{resource_id}/search_for_vaddr", self.search_for_vaddr),
                 web.post("/{resource_id}/add_tag", self.add_tag),
+                web.post(
+                    "/{resource_id}/add_flush_to_disk_to_script", self.add_flush_to_disk_to_script
+                ),
                 web.get("/get_all_tags", self.get_all_tags),
                 web.get("/{resource_id}/get_script", self.get_script),
                 web.get("/", self.get_static_files),
@@ -651,6 +654,22 @@ async def get_all_tags(self, request: Request) -> Response:
             self._serializer.to_pjson(self._ofrak_context.get_all_tags(), Set[ResourceTag])
         )
 
+    @exceptions_to_http(SerializedError)
+    async def add_flush_to_disk_to_script(self, request: Request) -> Response:
+        resource = await self._get_resource_for_request(request)
+        output_file_name = self._serializer.from_pjson(await request.json(), str)
+        # Use FilesystemEntry name as filename if available, otherwise generate random filename
+        if not output_file_name:
+            output_file_name = self._ofrak_context.id_service.generate_id().hex()
+        script_str = (
+            """
+        await {resource}"""
+            f""".flush_to_disk("{output_file_name}")"""
+        )
+        await self.script_builder.add_action(resource, script_str, ActionType.UNDEF)
+        await self.script_builder.commit_to_script(resource)
+        return json_response(self._serialize_resource(resource))
+
     @exceptions_to_http(SerializedError)
     async def get_script(self, request: Request) -> Response:
         resource = await self._get_resource_for_request(request)

ofrak_core/test_ofrak/unit/test_ofrak_server.py

@@ -1,5 +1,6 @@
 import itertools
 import json
+import os
 import pytest
 import sys
 
@@ -21,6 +22,16 @@ def hello_world_elf() -> bytes:
     return hello_elf()
 
 
+@pytest.fixture(scope="session")
+def firmware_zip() -> bytes:
+    assets_dir = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "../components/assets/binwalk_assets")
+    )
+    asset_path = os.path.join(assets_dir, "firmware.zip")
+    with open(asset_path, "rb") as f:
+        return f.read()
+
+
 # Create test server that will be spun up for each test
 @pytest.fixture
 async def ofrak_server(ofrak, ofrak_context):
@@ -666,3 +677,92 @@ async def test_clear_action_queue(ofrak_client: TestClient, hello_world_elf):
         "    ofrak.run(main)",
         "",
     ]
+
+
+async def test_add_flush_to_disk_to_script(ofrak_client: TestClient, firmware_zip):
+    create_resp = await ofrak_client.post(
+        "/create_root_resource", params={"name": "firmware_zip"}, data=firmware_zip
+    )
+    root = await create_resp.json()
+    root_id = root["id"]
+
+    await ofrak_client.post(f"/{root_id}/unpack")
+    child_resp = await ofrak_client.post(f"/batch/get_children", json=[root_id])
+    child_body = await child_resp.json()
+    only_child_id = child_body[root_id][0]["id"]
+    await ofrak_client.post(f"/{only_child_id}/unpack")
+    grandchildren_resp = await ofrak_client.post(f"/batch/get_children", json=[only_child_id])
+    grandchildren_body = await grandchildren_resp.json()
+    eldest_grandchild_id = grandchildren_body[only_child_id][0]["id"]
+    await ofrak_client.post(
+        f"/{eldest_grandchild_id}/add_flush_to_disk_to_script", json="DIR655B1_FW203NAB02.bin"
+    )
+
+    resp = await ofrak_client.get(
+        f"/{root_id}/get_script",
+    )
+    resp_body = await resp.json()
+    assert resp_body == [
+        "from ofrak import *",
+        "from ofrak.core import *",
+        "",
+        "",
+        "async def main(ofrak_context: OFRAKContext):",
+        "",
+        '    root_resource = await ofrak_context.create_root_resource_from_file("firmware_zip")',
+        "",
+        "    await root_resource.unpack()",
+        "",
+        "    folder_dir655_revB_FW_203NA = await root_resource.get_only_child(",
+        "        r_filter=ResourceFilter(",
+        "            tags={Folder},",
+        "            attribute_filters=[",
+        "                ResourceAttributeValueFilter(",
+        "                    attribute=AttributesType[FilesystemEntry].Name,",
+        '                    value="dir655_revB_FW_203NA",',
+        "                )",
+        "            ],",
+        "        )",
+        "    )",
+        "",
+        "    await folder_dir655_revB_FW_203NA.unpack()",
+        "",
+        "    file_DIR655B1_FW203NAB02_bin = await folder_dir655_revB_FW_203NA.get_only_child(",
+        "        r_filter=ResourceFilter(",
+        "            tags={File},",
+        "            attribute_filters=[",
+        "                ResourceAttributeValueFilter(",
+        "                    attribute=AttributesType[FilesystemEntry].Name,",
+        '                    value="DIR655B1_FW203NAB02.bin",',
+        "                )",
+        "            ],",
+        "        )",
+        "    )",
+        "",
+        '    await file_DIR655B1_FW203NAB02_bin.flush_to_disk("DIR655B1_FW203NAB02.bin")',
+        "",
+        "",
+        'if __name__ == "__main__":',
+        "    ofrak = OFRAK()",
+        "    if False:",
+        "        import ofrak_angr",
+        "        import ofrak_capstone",
+        "",
+        "        ofrak.discover(ofrak_capstone)",
+        "        ofrak.discover(ofrak_angr)",
+        "",
+        "    if False:",
+        "        import ofrak_binary_ninja",
+        "        import ofrak_capstone",
+        "",
+        "        ofrak.discover(ofrak_capstone)",
+        "        ofrak.discover(ofrak_binary_ninja)",
+        "",
+        "    if False:",
+        "        import ofrak_ghidra",
+        "",
+        "        ofrak.discover(ofrak_ghidra)",
+        "",
+        "    ofrak.run(main)",
+        "",
+    ]