diff --git a/frontend/public/icons/back-arrow.svg b/frontend/public/icons/back-arrow.svg new file mode 100644 index 000000000..b403baa1e --- /dev/null +++ b/frontend/public/icons/back-arrow.svg @@ -0,0 +1,35 @@ + + + + + + + + + + diff --git a/frontend/public/icons/binary.svg b/frontend/public/icons/binary.svg new file mode 100644 index 000000000..9bc61da1e --- /dev/null +++ b/frontend/public/icons/binary.svg @@ -0,0 +1,4 @@ + + + + diff --git a/frontend/public/icons/briefcase.svg b/frontend/public/icons/briefcase.svg new file mode 100644 index 000000000..75fd6ad12 --- /dev/null +++ b/frontend/public/icons/briefcase.svg @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/frontend/public/icons/disk.svg b/frontend/public/icons/disk.svg new file mode 100644 index 000000000..a302d568d --- /dev/null +++ b/frontend/public/icons/disk.svg @@ -0,0 +1,18 @@ + + + + + + + + + + + diff --git a/frontend/public/icons/reset.svg b/frontend/public/icons/reset.svg new file mode 100644 index 000000000..2ef9e199a --- /dev/null +++ b/frontend/public/icons/reset.svg @@ -0,0 +1,4 @@ + + + + diff --git a/frontend/public/icons/trash.svg b/frontend/public/icons/trash.svg new file mode 100644 index 000000000..fda1329d4 --- /dev/null +++ b/frontend/public/icons/trash.svg @@ -0,0 +1,23 @@ + + + + + + + + + + + diff --git a/frontend/src/App.svelte b/frontend/src/App.svelte index a8642fe4e..c5ba22244 100644 --- a/frontend/src/App.svelte +++ b/frontend/src/App.svelte @@ -51,10 +51,12 @@ import { keyEventToString, shortcuts } from "./keyboard.js"; import { writable } from "svelte/store"; + import ProjectManagerView from "./ProjectManager/ProjectManagerView.svelte"; printConsoleArt(); let showRootResource = false, + showProjectManager = false, dataLenPromise = Promise.resolve([]), hexScrollY = writable({}), useAssemblyView = false, @@ -188,6 +190,8 @@ Answer by running riddle.answer('your answer here') from the console.`); bind:bottomLeftPane="{bottomLeftPane}" bind:resourceNodeDataMap="{resourceNodeDataMap}" bind:modifierView="{modifierView}" + bind:showProjectManager="{showProjectManager}" + bind:showRootResource="{showRootResource}" /> {/if} @@ -251,10 +255,19 @@ Answer by running riddle.answer('your answer here') from the console.`); {/if} +{:else if showProjectManager} + {:else} + import Checkbox from "./Checkbox.svelte"; + + export let selectedValue, + ownValue, + leftbox = false, + nomargin = false; + + let thisSelected; + + $: if (thisSelected) { + selectedValue = ownValue; + } + + + + + diff --git a/frontend/src/FileBrowser.svelte b/frontend/src/FileBrowser.svelte index afd6ff96a..c987bbef7 100644 --- a/frontend/src/FileBrowser.svelte +++ b/frontend/src/FileBrowser.svelte @@ -31,7 +31,6 @@ .filelabel span { width: 100%; - margin-left: 2ch; background: inherit; color: inherit; /* border-bottom: 1px solid var(--main-fg-color); */ diff --git a/frontend/src/ProjectManager/ProjectManagerAddBinaryToProject.svelte b/frontend/src/ProjectManager/ProjectManagerAddBinaryToProject.svelte new file mode 100644 index 000000000..9ea48faa9 --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerAddBinaryToProject.svelte @@ -0,0 +1,79 @@ + + + + +
+ + {#if f} + + {/if} +
diff --git a/frontend/src/ProjectManager/ProjectManagerAddScriptToProject.svelte b/frontend/src/ProjectManager/ProjectManagerAddScriptToProject.svelte new file mode 100644 index 000000000..f4c7b88cb --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerAddScriptToProject.svelte @@ -0,0 +1,69 @@ + + + + +
+ + {#if f} + + {/if} +
diff --git a/frontend/src/ProjectManager/ProjectManagerBinaryOptions.svelte b/frontend/src/ProjectManager/ProjectManagerBinaryOptions.svelte new file mode 100644 index 000000000..e75ebaacf --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerBinaryOptions.svelte @@ -0,0 +1,68 @@ + + + + +
+ +
diff --git a/frontend/src/ProjectManager/ProjectManagerCheckbox.svelte b/frontend/src/ProjectManager/ProjectManagerCheckbox.svelte new file mode 100644 index 000000000..223f9ba46 --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerCheckbox.svelte @@ -0,0 +1,104 @@ + + + + +
+ + {#if inclusiveSelectionGroup !== undefined} + + + + {/if} + {#if exclusiveSelectionValue !== undefined} + + + + {/if} + +
diff --git a/frontend/src/ProjectManager/ProjectManagerFocusableLabel.svelte b/frontend/src/ProjectManager/ProjectManagerFocusableLabel.svelte new file mode 100644 index 000000000..69f758d76 --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerFocusableLabel.svelte @@ -0,0 +1,18 @@ + + + + + diff --git a/frontend/src/ProjectManager/ProjectManagerMainOptions.svelte b/frontend/src/ProjectManager/ProjectManagerMainOptions.svelte new file mode 100644 index 000000000..1e2efefb9 --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerMainOptions.svelte @@ -0,0 +1,9 @@ + + + + +
+

Welcome to the OFRAK Project Manager

+
diff --git a/frontend/src/ProjectManager/ProjectManagerOptions.svelte b/frontend/src/ProjectManager/ProjectManagerOptions.svelte new file mode 100644 index 000000000..e503844fd --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerOptions.svelte @@ -0,0 +1,28 @@ + + + + +
+
+ {#if focus && typeof focus == "string"} + "{focus}" + {:else if focus && typeof focus == "object"} + + {:else} + Click anywhere to see its options. + {/if} +
+
diff --git a/frontend/src/ProjectManager/ProjectManagerScriptOptions.svelte b/frontend/src/ProjectManager/ProjectManagerScriptOptions.svelte new file mode 100644 index 000000000..e616ada8d --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerScriptOptions.svelte @@ -0,0 +1,68 @@ + + + + +
+ +
diff --git a/frontend/src/ProjectManager/ProjectManagerToolbar.svelte b/frontend/src/ProjectManager/ProjectManagerToolbar.svelte new file mode 100644 index 000000000..28700ba9b --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerToolbar.svelte @@ -0,0 +1,82 @@ + + + diff --git a/frontend/src/ProjectManager/ProjectManagerView.svelte b/frontend/src/ProjectManager/ProjectManagerView.svelte new file mode 100644 index 000000000..f2ea11c31 --- /dev/null +++ b/frontend/src/ProjectManager/ProjectManagerView.svelte @@ -0,0 +1,226 @@ + + + + +
OFRAK Project Manager
+
+ +
+ + + +
+ +
+
+
+ {#each binariesForProject as binaryName} +
+ +
+ {/each} +
+
+
+ +
+ +
+
+
+
+ {#if scriptCheckboxHoverInfo.onInclusive} +

+ Script is {#if !scriptCheckboxHoverInfo.inclusiveChecked} + not + {/if} compatible with this binary +

+ {:else if scriptCheckboxHoverInfo.onExclusive} +

+ Script is {#if !scriptCheckboxHoverInfo.exclusiveChecked} + not + {/if} the one used to launch this binary +

+ {/if} +
+ {#each $selectedProject.scripts as script} +
+ {#if selectedBinaryName} + {#key forceRefreshProject} + + {/key} + {:else} + + {/if} +
+ {/each} +
+
+
+
+ + + +
+
+
diff --git a/frontend/src/ResourceTreeToolbar.svelte b/frontend/src/ResourceTreeToolbar.svelte index 3e43c4fc5..0d322d830 100644 --- a/frontend/src/ResourceTreeToolbar.svelte +++ b/frontend/src/ResourceTreeToolbar.svelte @@ -7,12 +7,21 @@ import SettingsView from "./SettingsView.svelte"; import Toolbar from "./Toolbar.svelte"; - import { selectedResource, selected, settings } from "./stores.js"; + import { + selectedResource, + selected, + settings, + selectedProject, + } from "./stores.js"; import SearchView from "./SearchView.svelte"; import AddTagView from "./AddTagView.svelte"; import RunScriptView from "./RunScriptView.svelte"; - export let resourceNodeDataMap, modifierView, bottomLeftPane; + export let resourceNodeDataMap, + modifierView, + bottomLeftPane, + showProjectManager, + showRootResource; $: rootResource = $selectedResource; function refreshResource() { @@ -22,7 +31,7 @@ $selected = originalSelected; } - let toolbarButtons, experimentalFeatures; + let toolbarButtons, experimentalFeatures, projectFeatures; const neverResolves = new Promise(() => {}); $: { experimentalFeatures = [ @@ -44,6 +53,25 @@ ]; } + $: { + projectFeatures = [ + { + text: "Project Manager", + iconUrl: "/icons/briefcase.svg", + onclick: async (e) => { + if ($selectedProject) { + let state = { + $selectProject: $selectedProject, + }; + showProjectManager = true; + showRootResource = false; + history.pushState(state, "", "/"); + } + }, + }, + ]; + } + $: { toolbarButtons = [ { @@ -248,6 +276,21 @@ }, }, + { + text: "Project Manager", + iconUrl: "/icons/briefcase.svg", + onclick: async (e) => { + if ($selectedProject) { + let state = { + $selectProject: $selectedProject, + }; + showProjectManager = true; + showRootResource = false; + history.pushState(state, "", "/"); + } + }, + }, + { text: "Search", iconUrl: "/icons/identify.svg", @@ -274,6 +317,10 @@ ]; } + $: if ($settings.$selectedProject) { + toolbarButtons = [...toolbarButtons, ...projectFeatures]; + } + $: if ($settings.experimentalFeatures) { toolbarButtons = [...toolbarButtons, ...experimentalFeatures]; } diff --git a/frontend/src/ResourceTreeView.svelte b/frontend/src/ResourceTreeView.svelte index a0acaf1f2..2bbc89846 100644 --- a/frontend/src/ResourceTreeView.svelte +++ b/frontend/src/ResourceTreeView.svelte @@ -53,7 +53,9 @@ export let rootResource, modifierView, bottomLeftPane, - resourceNodeDataMap = {}; + resourceNodeDataMap = {}, + showProjectManager, + showRootResource; let searchFilter; let searchResults = {}; @@ -78,6 +80,8 @@ bind:resourceNodeDataMap="{resourceNodeDataMap}" bind:modifierView="{modifierView}" bind:bottomLeftPane="{bottomLeftPane}" + bind:showProjectManager="{showProjectManager}" + bind:showRootResource="{showRootResource}" />
diff --git a/frontend/src/RunScriptView.svelte b/frontend/src/RunScriptView.svelte index 04d4039de..f1a67377a 100644 --- a/frontend/src/RunScriptView.svelte +++ b/frontend/src/RunScriptView.svelte @@ -1,9 +1,19 @@ @@ -310,34 +461,43 @@ on:drop|preventDefault="{handleDrop}" on:mousemove="{(e) => (mouseX = e.clientX)}" on:mouseleave="{() => (mouseX = undefined)}" - on:click="{() => fileinput.click()}" + on:click="{() => { + if (!showProjectOptions) { + fileinput.click(); + } + }}" style:border-color="{animals[selectedAnimal]?.color || "var(--main-fg-color)"}" style:color="{animals[selectedAnimal]?.color || "var(--main-fg-color)"}" > - {#if !dragging} + {#if !dragging && !showProjectOptions}

Drag in a file to analyze

Click anwyhere to browse for a file to analyze

- {:else} + {:else if dragging}

Drop the file!

- {/if} - - - -
+ {:else if showProjectOptions} +

Project Options

- OR - -
- + /> + {/if} + {#if !showProjectOptions} + + +
+ + OR + +
+ {/if} {#await preExistingRootsPromise} {:then preExistingRootResources} - {#if preExistingRootsPromise && preExistingRootsPromise.length > 0} + {#if !showProjectOptions && preExistingRootsPromise && preExistingRootsPromise.length > 0}
+ +
+ + OR + +
+ {#await preExistingProjectsPromise then projects} + + + {/await} +
+ + OR + +
+ + +
+ +
+
+ Show Advanced Options +
+ {#if showAdvancedProjectOptions} +
+ + +
+ {/if} +
+ + {:else} + + {/if} + + {/if} Response: return json_response(results) + @exceptions_to_http(SerializedError) + async def create_new_project(self, request: Request) -> Response: + if self.projects is None: + self.projects = self._slurp_projects_from_dir() + body = await request.json() + name = body.get("name") + project = OfrakProject.create(name, os.path.join(self.projects_dir, name)) + self.projects.add(project) + + return json_response({"id": project.session_id.hex()}) + + @exceptions_to_http(SerializedError) + async def clone_project_from_git(self, request: Request) -> Response: + if self.projects is None: + self.projects = self._slurp_projects_from_dir() + + def recurse_path_collisions(path: str, count: int) -> str: + if count == 0: + incr_path = path + else: + incr_path = f"{path}_{count}" + if os.path.exists(incr_path): + count += 1 + return recurse_path_collisions(path, count) + else: + return incr_path + + body = await request.json() + url = body.get("url") + path = recurse_path_collisions( + os.path.join(self.projects_dir, url.split(":")[-1].split("/")[-1]), 0 + ) + project = OfrakProject.clone_from_git(url, path) + self.projects.add(project) + return json_response({"id": project.session_id.hex()}) + + @exceptions_to_http(SerializedError) + async def get_project_by_id(self, request: Request) -> Response: + id = request.query.get("id") + project = self._get_project_by_id(id) + return json_response(project.get_current_metadata()) + + @exceptions_to_http(SerializedError) + async def get_all_projects(self, request: Request) -> Response: + if self.projects is None: + self.projects = self._slurp_projects_from_dir() + return json_response([project.get_current_metadata() for project in self.projects]) + + @exceptions_to_http(SerializedError) + async def reset_project(self, request: Request) -> Response: + body = await request.json() + id = body["id"] + project = self._get_project_by_id(id) + project.reset_project() + return json_response([]) + + @exceptions_to_http(SerializedError) + async def add_binary_to_project(self, request: Request) -> Response: + id = request.query.get("id") + name_query = request.query.get("name") + if name_query is not None: + name = name_query + data = await request.read() + project = self._get_project_by_id(id) + project.add_binary(name, data) + return json_response([]) + + @exceptions_to_http(SerializedError) + async def add_script_to_project(self, request: Request) -> Response: + id = request.query.get("id") + name_query = request.query.get("name") + if name_query is not None: + name = name_query + data = await request.read() + project = self._get_project_by_id(id) + project.add_script(name, data.decode()) + return json_response([]) + + @exceptions_to_http(SerializedError) + async def open_project(self, request: Request) -> Response: + body = await request.json() + id = body["id"] + binary = body["binary"] + script = body["script"] + if request.remote is not None: + resource_id = request.remote + else: + raise AttributeError("No resource ID provided") + project = self._get_project_by_id(id) + resource = await project.init_project_binary( + binary, self._ofrak_context, script_name=script + ) + self._job_ids[resource_id] = resource.get_job_id() + return json_response(self._serialize_resource(resource)) + + @exceptions_to_http(SerializedError) + async def get_projects_path(self, request: Request) -> Response: + return json_response(self.projects_dir) + + @exceptions_to_http(SerializedError) + async def set_projects_path(self, request: Request) -> Response: + body = await request.json() + new_path = body["path"] + if not os.path.exists(new_path): + os.mkdir(new_path) + self.projects_dir = new_path + self.projects = self._slurp_projects_from_dir() + return json_response(self.projects_dir) + + @exceptions_to_http(SerializedError) + async def update_binary_data(self, request: Request) -> Response: + body = await request.json() + id = body["id"] + binary_name = body["name"] + init_script = body["init"] + associated_scripts = body["associated_scripts"] + project = self._get_project_by_id(id) + project.update_binary_data(binary_name, init_script, associated_scripts) + return json_response([]) + + @exceptions_to_http(SerializedError) + async def save_project_data(self, request: Request) -> Response: + body = await request.json() + session_id = body["session_id"] + project = self._get_project_by_id(session_id) + for binary_name, binary_metadata in body["binaries"].items(): + project.binaries[binary_name].init_script = binary_metadata["init_script"] + project.binaries[binary_name].associated_scripts = binary_metadata["associated_scripts"] + project.write_metadata_to_disk() + return json_response([]) + + @exceptions_to_http(SerializedError) + async def delete_binary_from_project(self, request: Request) -> Response: + body = await request.json() + id = body["id"] + binary_name = body["binary"] + project = self._get_project_by_id(id) + project.delete_binary(binary_name) + return json_response([]) + + @exceptions_to_http(SerializedError) + async def delete_script_from_project(self, request: Request) -> Response: + body = await request.json() + id = body["id"] + script_name = body["script"] + project = self._get_project_by_id(id) + project.delete_script(script_name) + return json_response([]) + + @exceptions_to_http(SerializedError) + async def get_project_script(self, request: Request) -> Response: + project_id = request.query.get("project") + script_name_query = request.query.get("script") + if script_name_query is not None: + script_name = script_name_query + project = self._get_project_by_id(project_id) + script_body = project.get_script_body(script_name) + + return Response(text=script_body) + + def _slurp_projects_from_dir(self) -> Set: + projects = set() + if not os.path.exists(self.projects_dir): + os.makedirs(self.projects_dir) + for dir in os.listdir(self.projects_dir): + try: + project = OfrakProject.init_from_path(os.path.join(self.projects_dir, dir)) + projects.add(project) + except: + pass + return projects + + def _get_project_by_name(self, name) -> Optional[OfrakProject]: + if self.projects is None: + self.projects = self._slurp_projects_from_dir() + result = [project for project in self.projects if project.name == name] + if len(result) > 1: + raise AttributeError("Project Name Collision") + if len(result) == 0: + return None + return result[0] + + def _get_project_by_id(self, id) -> OfrakProject: + if self.projects is None: + self.projects = self._slurp_projects_from_dir() + result = [project for project in self.projects if project.session_id.hex() == id] + if len(result) > 1: + raise AttributeError("Project ID Collision") + if len(result) == 0: + raise ValueError(f"Project with ID {id} not found") + return result[0] + def _construct_field_response(self, obj): if dataclasses.is_dataclass(obj): res = [] diff --git a/ofrak_core/ofrak/project/__init__.py b/ofrak_core/ofrak/project/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/ofrak_core/ofrak/project/project.py b/ofrak_core/ofrak/project/project.py new file mode 100644 index 000000000..312c29894 --- /dev/null +++ b/ofrak_core/ofrak/project/project.py @@ -0,0 +1,382 @@ +import binascii +import json +import os.path +import uuid +from git import Repo +from dataclasses import dataclass +from typing import Dict, List, Optional + +from ofrak.core.run_script_modifier import RunScriptModifier, RunScriptModifierConfig + +from ofrak.resource import Resource + +from ofrak.ofrak_context import OFRAKContext + + +@dataclass +class _OfrakProjectBinary: + associated_scripts: List[str] + init_script: Optional[str] + + +class OfrakProject: + """ + An OFRAK 'project' + + """ + + def __init__( + self, + path: str, + name: str, + project_id: bytes, + binaries: Dict[str, _OfrakProjectBinary], + scripts: List[str], + ): + self.path: str = path + self.name: str = name + self.project_id: bytes = project_id + self.binaries: Dict[str, _OfrakProjectBinary] = binaries + self.scripts: List[str] = scripts + self.session_id = uuid.uuid4().bytes + + @property + def metadata_path(self): + return os.path.join(self.path, "metadata.json") + + @property + def readme_path(self): + return os.path.join(self.path, "README.md") + + @staticmethod + def create(name: str, path: str) -> "OfrakProject": + new_project = OfrakProject( + path, + name, + uuid.uuid4().bytes, + {}, + [], + ) + + os.makedirs(os.path.join(path, "scripts"), exist_ok=True) + os.makedirs(os.path.join(path, "binaries"), exist_ok=True) + + new_project.write_metadata_to_disk() + + return new_project + + @staticmethod + def clone_from_git(url: str, path: str) -> "OfrakProject": + repo = Repo.clone_from(url, path) + if not os.path.exists(path): + raise ValueError(f"{path} does not exist") + if not os.path.isdir(path): + raise ValueError(f"{path} is not a directory") + + metadata_path = os.path.join(path, "metadata.json") + binaries_path = os.path.join(path, "binaries") + scripts_path = os.path.join(path, "scripts") + + if not all( + [ + os.path.exists(metadata_path), + os.path.exists(binaries_path), + os.path.isdir(binaries_path), + os.path.exists(scripts_path), + os.path.isdir(scripts_path), + ] + ): + raise ValueError(f"{path} has invalid structure to be an Project") + + with open(metadata_path) as f: + raw_metadata = json.load(f) + + scripts = [script["name"] for script in raw_metadata["scripts"]] + + binaries = {} + + for binary_name, info in raw_metadata["binaries"].items(): + binaries[binary_name] = _OfrakProjectBinary( + info["associated_scripts"], info.get("init_script") + ) + name = raw_metadata["name"] + project_id = binascii.unhexlify(raw_metadata["project_id"]) + project = OfrakProject( + path, + name, + project_id, + binaries, + scripts, + ) + + return project + + @staticmethod + def init_from_path(path: str) -> "OfrakProject": + """ + + Assume path points to a directory with the following structure: + (top-level directory) + |-metadata.json + |-README.md + |--binaries + | |-binary1.bin + | | ... + |--scripts + |-script1.py + | ... + + :param path: + :return: + """ + if not os.path.exists(path): + raise ValueError(f"{path} does not exist") + if not os.path.isdir(path): + raise ValueError(f"{path} is not a directory") + + metadata_path = os.path.join(path, "metadata.json") + binaries_path = os.path.join(path, "binaries") + scripts_path = os.path.join(path, "scripts") + + if not all( + [ + os.path.exists(metadata_path), + os.path.exists(binaries_path), + os.path.isdir(binaries_path), + os.path.exists(scripts_path), + os.path.isdir(scripts_path), + ] + ): + raise ValueError(f"{path} has invalid structure to be an Project") + + with open(metadata_path) as f: + raw_metadata = json.load(f) + + scripts = [script["name"] for script in raw_metadata["scripts"]] + + binaries = {} + + for binary_name, info in raw_metadata["binaries"].items(): + binaries[binary_name] = _OfrakProjectBinary( + info["associated_scripts"], info.get("init_script") + ) + name = raw_metadata["name"] + project_id = binascii.unhexlify(raw_metadata["project_id"]) + project = OfrakProject( + path, + name, + project_id, + binaries, + scripts, + ) + + return project + + def script_path(self, script_name, check: bool = True) -> str: + if check and script_name not in self.scripts: + raise ValueError(f"Script {script_name} is not a script in this Project") + p = os.path.join(self.path, "scripts", script_name) + if check and not os.path.exists(p): + raise ValueError( + f"Script {script_name} is known to this Project but is not on disk " + f"(looked at {p})" + ) + return p + + def binary_path(self, binary_name, check: bool = True) -> str: + if check and binary_name not in self.binaries: + raise ValueError(f"Binary {binary_name} is not a binary in this Project") + p = os.path.join(self.path, "binaries", binary_name) + if check and not os.path.exists(p): + raise ValueError( + f"Binary {binary_name} is known to this Project but is not on disk " + f"(looked at {p})" + ) + return p + + async def init_project_binary( + self, binary_name: str, ofrak_context: OFRAKContext, script_name: Optional[str] = None + ) -> Resource: + if script_name is None: + binary_metadata = self.binaries[binary_name] + script_name = binary_metadata.init_script + + resource = await ofrak_context.create_root_resource_from_file(self.binary_path(binary_name)) + + if script_name is not None: + with open(self.script_path(script_name)) as f: + code = f.read() + await resource.run(RunScriptModifier, RunScriptModifierConfig(code)) + + return resource + + def get_script_body(self, script_name: str) -> str: + if script_name not in self.scripts: + raise ValueError(f"Script {script_name} is not a script in this Project") + with open(self.script_path(script_name)) as f: + code = f.read() + return code + + def write_metadata_to_disk(self): + metadata = { + "name": self.name, + "project_id": self.project_id.hex(), + "scripts": [ + { + "name": script_name, + } + for script_name in self.scripts + ], + "binaries": { + binary_name: { + "init_script": binary_info.init_script, + "associated_scripts": binary_info.associated_scripts, + } + for binary_name, binary_info in self.binaries.items() + }, + } + with open(os.path.join(self.path, "metadata.json"), "w") as f: + json.dump(metadata, f) + + def get_saved_metadata(self): + with open(os.path.join(self.path, "metadata.json")) as f: + data = json.load(f) + data["session_id"] = self.session_id.hex() + return data + + def get_current_metadata(self): + return { + "name": self.name, + "project_id": self.project_id.hex(), + "session_id": self.session_id.hex(), + "scripts": [ + { + "name": script_name, + } + for script_name in self.scripts + ], + "binaries": { + binary_name: { + "init_script": binary_info.init_script, + "associated_scripts": binary_info.associated_scripts, + } + for binary_name, binary_info in self.binaries.items() + }, + } + + def add_binary(self, name: str, contents: bytes): + self.binaries[name] = _OfrakProjectBinary([], None) + os.makedirs(os.path.join(self.path, "binaries"), exist_ok=True) + with open(self.binary_path(name, check=False), "wb+") as f: + f.write(contents) + + def add_script(self, name: str, script_contents: str): + self.scripts.append(name) + os.makedirs(os.path.join(self.path, "scripts"), exist_ok=True) + with open(self.script_path(name, check=False), "w+") as f: + f.write(script_contents) + + def update_binary_data( + self, name: str, init: Optional[str] = None, associated: Optional[List[str]] = None + ): + binary = self._get_binary(name) + if init is not None: + binary.init_script = init + if associated is not None: + binary.associated_scripts = associated + + def delete_binary(self, name: str): + if not os.path.isdir(os.path.join(self.path, ".Trash")): + os.mkdir(os.path.join(self.path, ".Trash")) + if not os.path.isdir(os.path.join(os.path.join(self.path, ".Trash"), "binaries")): + os.mkdir(os.path.join(os.path.join(self.path, ".Trash"), "binaries")) + os.rename( + self.binary_path(name), + os.path.join(os.path.join(os.path.join(self.path, ".Trash"), "binaries"), name), + ) + self.binaries.pop(name) + + def delete_script(self, name: str): + if not os.path.isdir(os.path.join(self.path, ".Trash")): + os.mkdir(os.path.join(self.path, ".Trash")) + if not os.path.isdir(os.path.join(os.path.join(self.path, ".Trash"), "scripts")): + os.mkdir(os.path.join(os.path.join(self.path, ".Trash"), "scripts")) + os.rename( + self.script_path(name), + os.path.join(os.path.join(os.path.join(self.path, ".Trash"), "scripts"), name), + ) + self.scripts.remove(name) + + def reset_project(self): + path = self.path + if not os.path.exists(path): + raise ValueError(f"{path} does not exist") + if not os.path.isdir(path): + raise ValueError(f"{path} is not a directory") + + metadata_path = os.path.join(path, "metadata.json") + binaries_path = os.path.join(path, "binaries") + scripts_path = os.path.join(path, "scripts") + + if not all( + [ + os.path.exists(metadata_path), + os.path.exists(binaries_path), + os.path.isdir(binaries_path), + os.path.exists(scripts_path), + os.path.isdir(scripts_path), + ] + ): + raise ValueError(f"{path} has invalid structure to be an Project") + + with open(metadata_path) as f: + raw_metadata = json.load(f) + + self.scripts = [script["name"] for script in raw_metadata["scripts"]] + for script in self.scripts: + if not os.path.exists(self.script_path(script, check=False)): + if not os.path.exists( + os.path.join(os.path.join(os.path.join(self.path, ".Trash"), "scripts"), script) + ): + raise AttributeError( + f"Trying to restore script {script} but the file is missing from .Trash" + ) + else: + os.rename( + os.path.join( + os.path.join(os.path.join(self.path, ".Trash"), "scripts"), script + ), + self.script_path(script, check=False), + ) + + self.binaries = {} + + for binaryName, info in raw_metadata["binaries"].items(): + self.binaries[binaryName] = _OfrakProjectBinary( + info["associated_scripts"], info.get("init_script") + ) + for binary in self.binaries.keys(): + if not os.path.exists(self.binary_path(binary, check=False)): + if not os.path.exists( + os.path.join( + os.path.join(os.path.join(self.path, ".Trash"), "binaries"), binary + ) + ): + raise AttributeError( + f"Trying to restore binary {binary} but the file is missing from .Trash" + ) + else: + os.rename( + os.path.join( + os.path.join(os.path.join(self.path, ".Trash"), "scripts"), binary + ), + self.binary_path(binary, check=False), + ) + + self.name = raw_metadata["name"] + self.project_id = binascii.unhexlify(raw_metadata["project_id"]) + + def _get_binary(self, name): + if not name in self.binaries.keys(): + raise ValueError(f"Binary with the name {name} does not exist") + return self.binaries[name] diff --git a/ofrak_core/test_ofrak/unit/ofrak_project/test_ofrak_project.py b/ofrak_core/test_ofrak/unit/ofrak_project/test_ofrak_project.py new file mode 100644 index 000000000..ed390d6ac --- /dev/null +++ b/ofrak_core/test_ofrak/unit/ofrak_project/test_ofrak_project.py @@ -0,0 +1,27 @@ +import os.path + +from ofrak import OFRAKContext +from ofrak.project.project import OfrakProject + +TEST_PROJECT_PATH = os.path.join(os.path.dirname(__file__), "test_projects", "project1") + + +async def test_load_project(ofrak_context: OFRAKContext): + project = OfrakProject.init_from_path(TEST_PROJECT_PATH) + initialized_resource = await project.init_project_binary("hello_world.bin", ofrak_context) + assert len(list(await initialized_resource.get_children())) > 0 + + +async def test_create_new_project(ofrak_context, tmpdir): + new_project = OfrakProject.create( + "New Test Project", + tmpdir, + ) + + new_project.add_binary("tiny_binary", b"just a basic binary, nothing special") + + new_project.write_metadata_to_disk() + + new_project2 = OfrakProject.init_from_path(tmpdir) + assert new_project.project_id == new_project2.project_id + assert new_project.binaries == new_project2.binaries diff --git a/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/README.md b/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/README.md new file mode 100644 index 000000000..e69de29bb diff --git a/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/binaries/hello_world.bin b/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/binaries/hello_world.bin new file mode 100755 index 000000000..d0340b0a2 Binary files /dev/null and b/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/binaries/hello_world.bin differ diff --git a/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/metadata.json b/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/metadata.json new file mode 100644 index 000000000..1c0eadbc5 --- /dev/null +++ b/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/metadata.json @@ -0,0 +1,14 @@ +{ + "name": "Test project A", + "project_id": "38f71e0209a6cc", + "scripts": [ + {"name": "script_a.py"} + ], + "binaries": { + "hello_world.bin": + { + "associated_scripts": ["script_a.py"], + "init_script": "script_a.py" + } + } +} diff --git a/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/scripts/script_a.py b/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/scripts/script_a.py new file mode 100644 index 000000000..91a546b8e --- /dev/null +++ b/ofrak_core/test_ofrak/unit/ofrak_project/test_projects/project1/scripts/script_a.py @@ -0,0 +1,9 @@ +from ofrak import * +from ofrak.core import * + + +async def main(ofrak_context: OFRAKContext, root_resource: Optional[Resource]): + if root_resource is None: + raise ValueError() + + await root_resource.unpack() diff --git a/ofrak_core/test_ofrak/unit/test_ofrak_server.py b/ofrak_core/test_ofrak/unit/test_ofrak_server.py index da0de69ab..33b64bb6f 100644 --- a/ofrak_core/test_ofrak/unit/test_ofrak_server.py +++ b/ofrak_core/test_ofrak/unit/test_ofrak_server.py @@ -1,6 +1,7 @@ import itertools import json import os +import shutil import tempfile from ofrak.ofrak_context import OFRAKContext from ofrak.resource import Resource @@ -1197,3 +1198,241 @@ async def test_search_data(ofrak_client: TestClient, hello_world_elf): resp_body2 = await resp.json() assert resp.status == 200 assert resp_body1 == resp_body2 + + +async def test_create_new_project(ofrak_client: TestClient): + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + assert resp.status == 200 + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_get_project_by_id(ofrak_client: TestClient): + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + resp_body = await resp.json() + id = resp_body["id"] + + resp = await ofrak_client.get("/get_project_by_id", params={"id": id}) + assert resp.status == 200 + body = await resp.json() + assert list(body.keys()) == ["name", "project_id", "session_id", "scripts", "binaries"] + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_get_all_projects(ofrak_client: TestClient): + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test1"}, + ) + resp_body = await resp.json() + id1 = resp_body["id"] + + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test2"}, + ) + resp_body = await resp.json() + id2 = resp_body["id"] + + resp = await ofrak_client.get("/get_all_projects") + assert resp.status == 200 + body = await resp.json() + assert len(body) == 2 + assert "test1" in [project["name"] for project in body] + assert "test2" in [project["name"] for project in body] + assert id1 in [project["session_id"] for project in body] + assert id2 in [project["session_id"] for project in body] + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_reset_project(ofrak_client: TestClient): + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.post( + "/reset_project", + json={"id": id}, + ) + assert resp.status == 200 + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_add_binary_to_project(ofrak_client: TestClient, hello_world_elf): + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.post( + "/add_binary_to_project", params={"id": id, "name": "hello_world_elf"}, data=hello_world_elf + ) + assert resp.status == 200 + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_add_script_to_project(ofrak_client: TestClient): + script = b"async def main(ofrak_context: OFRAKContext, root_resource: Optional[Resource] = None):\n\tawait root_resource.unpack()" + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.post( + "/add_script_to_project", params={"id": id, "name": "unpack.py"}, data=script + ) + assert resp.status == 200 + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_get_projects_path(ofrak_client: TestClient): + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.get("/get_projects_path") + resp_body = await resp.json() + assert resp_body == "/tmp/test-ofrak-projects" + + +async def test_save_project_data(ofrak_client: TestClient, hello_world_elf): + script = b"async def main(ofrak_context: OFRAKContext, root_resource: Optional[Resource] = None):\n\tawait root_resource.unpack()" + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.post( + "/add_script_to_project", params={"id": id, "name": "unpack.py"}, data=script + ) + resp = await ofrak_client.post( + "/add_binary_to_project", params={"id": id, "name": "hello_world_elf"}, data=hello_world_elf + ) + resp = await ofrak_client.post("/save_project_data", json={"sesion_id": id}) + resp = await ofrak_client.post("/reset_project", json={"sesion_id": id}) + resp = await ofrak_client.get("/get_all_projects") + resp_body = await resp.json() + resp = await ofrak_client.post( + "/add_binary_to_project", params={"id": id, "name": "hello_world_elf"}, data=hello_world_elf + ) + assert len(resp_body) == 1 + assert resp_body[0]["scripts"] == [{"name": "unpack.py"}] + assert resp_body[0]["binaries"] == { + "hello_world_elf": {"init_script": None, "associated_scripts": []} + } + assert resp.status == 200 + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_delete_from_project(ofrak_client: TestClient, hello_world_elf): + script = b"async def main(ofrak_context: OFRAKContext, root_resource: Optional[Resource] = None):\n\tawait root_resource.unpack()" + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.post( + "/add_script_to_project", params={"id": id, "name": "unpack.py"}, data=script + ) + resp = await ofrak_client.post( + "/add_binary_to_project", params={"id": id, "name": "hello_world_elf"}, data=hello_world_elf + ) + resp = await ofrak_client.post("/save_project_data", json={"sesion_id": id}) + resp = await ofrak_client.post("/reset_project", json={"sesion_id": id}) + resp = await ofrak_client.get("/get_all_projects") + resp_body = await resp.json() + assert len(resp_body) == 1 + assert resp_body[0]["scripts"] == [{"name": "unpack.py"}] + assert resp_body[0]["binaries"] == { + "hello_world_elf": {"init_script": None, "associated_scripts": []} + } + resp = await ofrak_client.post( + "/delete_script_from_project", + json={"id": id, "script": "unpack.py"}, + ) + resp = await ofrak_client.post( + "/delete_binary_from_project", + json={"id": id, "binary": "hello_world_elf"}, + ) + resp = await ofrak_client.post("/save_project_data", json={"session_id": id}) + resp = await ofrak_client.post("/reset_project", json={"session_id": id}) + resp = await ofrak_client.get("/get_all_projects") + resp_body = await resp.json() + assert resp_body[0]["scripts"] == [] + assert resp_body[0]["binaries"] == {} + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_get_project_script(ofrak_client: TestClient): + script = b"async def main(ofrak_context: OFRAKContext, root_resource: Optional[Resource] = None):\n\tawait root_resource.unpack()" + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post( + "/create_new_project", + json={"name": "test"}, + ) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.post( + "/add_script_to_project", params={"id": id, "name": "unpack.py"}, data=script + ) + resp = await ofrak_client.get( + "/get_project_script", + params={"project": id, "script": "unpack.py"}, + ) + assert resp.status == 200 + resp_body = await resp.text() + assert resp_body == script.decode() + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_git_clone_project(ofrak_client: TestClient): + git_url = "https://github.com/redballoonsecurity/ofrak-project-example.git" + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post("/clone_project_from_git", json={"url": git_url}) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.get("/get_project_by_id", params={"id": id}) + resp_body = await resp.json() + assert resp_body["scripts"] == [ + {"name": "unpack-and-comment.py"}, + {"name": "unpack.py"}, + {"name": "modify.py"}, + ] + assert resp_body["binaries"] == { + "example_program": { + "init_script": "modify.py", + "associated_scripts": ["unpack-and-comment.py", "unpack.py", "modify.py"], + } + } + shutil.rmtree("/tmp/test-ofrak-projects") + + +async def test_open_project(ofrak_client: TestClient): + git_url = "https://github.com/redballoonsecurity/ofrak-project-example.git" + await ofrak_client.post("/set_projects_path", json={"path": "/tmp/test-ofrak-projects"}) + resp = await ofrak_client.post("/clone_project_from_git", json={"url": git_url}) + resp_body = await resp.json() + id = resp_body["id"] + resp = await ofrak_client.post( + "/open_project", + json={"id": id, "binary": "example_program", "script": "unpack-and-comment.py"}, + ) + resp_body = await resp.json() + assert resp_body["id"] == "00000001" + shutil.rmtree("/tmp/test-ofrak-projects")