From d153b66106cbb6900fb96a0566b8ca51d7327839 Mon Sep 17 00:00:00 2001 From: John Lugton Date: Fri, 20 Sep 2024 15:01:45 +0100 Subject: [PATCH 1/6] Add output_path parameter to *System and tidy up a bit --- varc_core/systems/__init__.py | 9 +++++---- varc_core/systems/base_system.py | 21 ++++++++------------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/varc_core/systems/__init__.py b/varc_core/systems/__init__.py index 4b59744..b1b0870 100644 --- a/varc_core/systems/__init__.py +++ b/varc_core/systems/__init__.py @@ -10,7 +10,8 @@ def acquire_system( include_memory: bool = True, include_open: bool = True, extract_dumps: bool = False, - yara_file: Optional[str] = None + yara_file: Optional[str] = None, + output_path: Optional[str] = None ) -> BaseSystem: """Returns the either a windows or linux system or osx system @@ -20,12 +21,12 @@ def acquire_system( logging.info(f"Operating System is: {platform}") if platform == "linux" or platform == "linux2": from varc_core.systems.linux import LinuxSystem - return LinuxSystem(include_memory, include_open, extract_dumps, yara_file) + return LinuxSystem(include_memory, include_open, extract_dumps, yara_file, output_path=output_path) elif platform == "darwin": from varc_core.systems.osx import OsxSystem - return OsxSystem(include_memory, include_open, extract_dumps) + return OsxSystem(include_memory, include_open, extract_dumps, output_path=output_path) elif platform == "win32": from varc_core.systems.windows import WindowsSystem - return WindowsSystem(include_memory, include_open, extract_dumps, yara_file) + return WindowsSystem(include_memory, include_open, extract_dumps, yara_file, output_path=output_path) else: raise MissingOperatingSystemInfo() \ No newline at end of file diff --git a/varc_core/systems/base_system.py b/varc_core/systems/base_system.py index 091bd28..cdb2dbd 100644 --- a/varc_core/systems/base_system.py +++ b/varc_core/systems/base_system.py @@ -52,7 +52,8 @@ def __init__( include_memory: bool = True, include_open: bool = True, extract_dumps: bool = False, - yara_file: Optional[str] = None + yara_file: Optional[str] = None, + output_path: Optional[str] = None ) -> None: self.todays_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S") logging.info(f'Acquiring system: {self.get_machine_name()}, at {self.todays_date}') @@ -66,11 +67,13 @@ def __init__( self.yara_file = yara_file self.yara_results: List[dict] = [] self.yara_hit_pids: List[int] = [] + self.output_path = output_path or os.path.join("", f"{self.get_machine_name()}-{self.timestamp}.zip") if self.process_name and self.process_id: raise ValueError( "Only one of Process name or Process ID (PID) can be used. Please re-run using one or the other.") - self.zip_path = self.acquire_volatile() + + self.acquire_volatile() if self.yara_file: if not _YARA_AVAILABLE: @@ -273,11 +276,9 @@ def take_screenshot(self) -> Optional[bytes]: logging.error("Unable to take screenshot") return None - def acquire_volatile(self, output_path: Optional[str] = None) -> str: + def acquire_volatile(self) -> None: """Acquire volatile data into a zip file This is called by all OS's - - :return: The filepath of the zip """ self.process_info = self.get_processes() self.network_log = self.get_network() @@ -290,12 +291,8 @@ def acquire_volatile(self, output_path: Optional[str] = None) -> str: screenshot_image = self.take_screenshot() else: screenshot_image = None - if not output_path: - output_path = os.path.join("", f"{self.get_machine_name()}-{self.timestamp}.zip") - # strip .zip if in filename as shutil appends to end - archive_out = output_path + ".zip" if not output_path.endswith(".zip") else output_path - self.output_path = output_path - with zipfile.ZipFile(archive_out, 'a', compression=zipfile.ZIP_DEFLATED) as zip_file: + + with zipfile.ZipFile(self.output_path, 'a', compression=zipfile.ZIP_DEFLATED) as zip_file: if screenshot_image: zip_file.writestr(f"{self.get_machine_name()}-{self.timestamp}.png", screenshot_image) for key, value in table_data.items(): @@ -319,8 +316,6 @@ def acquire_volatile(self, output_path: Optional[str] = None) -> str: except FileNotFoundError: logging.warning(f"Could not open {file_path} for reading") - return archive_out - def yara_scan(self) -> None: def yara_hit_callback(hit: dict) -> Any: From 7c665f3de3cc117658796c57f087cec635e0c818 Mon Sep 17 00:00:00 2001 From: John Lugton Date: Fri, 20 Sep 2024 15:06:47 +0100 Subject: [PATCH 2/6] standardize zip interface --- varc_core/systems/base_system.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/varc_core/systems/base_system.py b/varc_core/systems/base_system.py index cdb2dbd..db70e9c 100644 --- a/varc_core/systems/base_system.py +++ b/varc_core/systems/base_system.py @@ -296,12 +296,10 @@ def acquire_volatile(self) -> None: if screenshot_image: zip_file.writestr(f"{self.get_machine_name()}-{self.timestamp}.png", screenshot_image) for key, value in table_data.items(): - with zip_file.open(f"{key}.json", 'w') as json_file: - json_file.write(value.encode()) + zip_file.writestr(f"{key}.json", value.encode()) if self.network_log: logging.info("Adding Netstat Data") - with zip_file.open("netstat.log", 'w') as network_file: - network_file.write("\r\n".join(self.network_log).encode()) + zip_file.writestr("netstat.log", "\r\n".join(self.network_log).encode()) if self.include_open and self.dumped_files: for file_path in self.dumped_files: logging.info(f"Adding open file {file_path}") From e97f56b1df93c528a3b1b1d6880e3fc5128363fc Mon Sep 17 00:00:00 2001 From: John Lugton Date: Fri, 20 Sep 2024 15:29:32 +0100 Subject: [PATCH 3/6] Allow output az lz4 --- requirements.txt | 3 ++- varc_core/systems/base_system.py | 40 ++++++++++++++++++++++++++++---- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/requirements.txt b/requirements.txt index 0a87ab9..45b29c0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ tqdm==4.64.1 Pymem==1.8.5 python-magic==0.4.24 pyinstaller # dont set a version, not compatible -yara-python==4.3.1 \ No newline at end of file +yara-python==4.3.1 +lz4==4.3.3 \ No newline at end of file diff --git a/varc_core/systems/base_system.py b/varc_core/systems/base_system.py index db70e9c..1b5fc47 100644 --- a/varc_core/systems/base_system.py +++ b/varc_core/systems/base_system.py @@ -7,17 +7,20 @@ If it can't work cross-platform, put any platform specific code in the class that inherits this base e.g. In linux.py """ +import io import json import logging import os import os.path import socket +import tarfile import time import zipfile from base64 import b64encode from datetime import datetime from typing import Any, List, Optional +import lz4.frame import mss import psutil from tqdm import tqdm @@ -292,14 +295,14 @@ def acquire_volatile(self) -> None: else: screenshot_image = None - with zipfile.ZipFile(self.output_path, 'a', compression=zipfile.ZIP_DEFLATED) as zip_file: + with self._open_output() as output_file: if screenshot_image: - zip_file.writestr(f"{self.get_machine_name()}-{self.timestamp}.png", screenshot_image) + output_file.writestr(f"{self.get_machine_name()}-{self.timestamp}.png", screenshot_image) for key, value in table_data.items(): - zip_file.writestr(f"{key}.json", value.encode()) + output_file.writestr(f"{key}.json", value.encode()) if self.network_log: logging.info("Adding Netstat Data") - zip_file.writestr("netstat.log", "\r\n".join(self.network_log).encode()) + output_file.writestr("netstat.log", "\r\n".join(self.network_log).encode()) if self.include_open and self.dumped_files: for file_path in self.dumped_files: logging.info(f"Adding open file {file_path}") @@ -308,12 +311,17 @@ def acquire_volatile(self) -> None: logging.warning(f"Skipping file as too large {file_path}") else: try: - zip_file.write(file_path, strip_drive(f"./collected_files/{file_path}")) + output_file.write(file_path, strip_drive(f"./collected_files/{file_path}")) except PermissionError: logging.warn(f"Permission denied copying {file_path}") except FileNotFoundError: logging.warning(f"Could not open {file_path} for reading") + def _open_output(self) -> zipfile.ZipFile | "_TarLz4Wrapper": + if self.output_path.endswith('.tar.lz4'): + return _TarLz4Wrapper(self.output_path) + else: + return zipfile.ZipFile(self.output_path, 'a', compression=zipfile.ZIP_DEFLATED) def yara_scan(self) -> None: def yara_hit_callback(hit: dict) -> Any: @@ -350,3 +358,25 @@ def yara_hit_callback(hit: dict) -> Any: logging.info("YARA scan results written to yara_results.json in output archive.") else: logging.info("No YARA rules were triggered. Nothing will be written to the output archive.") + + +class _TarLz4Wrapper: + + def __init__(self, path) -> None: + self._lz4 = lz4.frame.open(path, 'wb') + self._tar = tarfile.open(fileobj=self._lz4f, mode="w") + + def writestr(self, path: str, value: str | bytes): + info = tarfile.TarInfo(path) + info.size = len(value) + self._tar.addfile(info, io.BytesIO(value)) + + def write(self, path: str, arcname: str): + self._tar.add(path, arcname) + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + self._tar.close() + self._lz4.close() \ No newline at end of file From 405a9b22d9370f1e8d2a99844f4d4c8dacf5a7c2 Mon Sep 17 00:00:00 2001 From: John Lugton Date: Fri, 20 Sep 2024 15:40:15 +0100 Subject: [PATCH 4/6] Fix type ref --- varc_core/systems/base_system.py | 45 ++++++++++++++++---------------- 1 file changed, 23 insertions(+), 22 deletions(-) diff --git a/varc_core/systems/base_system.py b/varc_core/systems/base_system.py index 1b5fc47..722bf51 100644 --- a/varc_core/systems/base_system.py +++ b/varc_core/systems/base_system.py @@ -36,6 +36,28 @@ _MAX_OPEN_FILE_SIZE = 10000000 # 10 Mb max dumped filesize +class _TarLz4Wrapper: + + def __init__(self, path) -> None: + self._lz4 = lz4.frame.open(path, 'wb') + self._tar = tarfile.open(fileobj=self._lz4f, mode="w") + + def writestr(self, path: str, value: str | bytes): + info = tarfile.TarInfo(path) + info.size = len(value) + self._tar.addfile(info, io.BytesIO(value)) + + def write(self, path: str, arcname: str): + self._tar.add(path, arcname) + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + self._tar.close() + self._lz4.close() + + class BaseSystem: """A @@ -317,7 +339,7 @@ def acquire_volatile(self) -> None: except FileNotFoundError: logging.warning(f"Could not open {file_path} for reading") - def _open_output(self) -> zipfile.ZipFile | "_TarLz4Wrapper": + def _open_output(self) -> zipfile.ZipFile | _TarLz4Wrapper: if self.output_path.endswith('.tar.lz4'): return _TarLz4Wrapper(self.output_path) else: @@ -359,24 +381,3 @@ def yara_hit_callback(hit: dict) -> Any: else: logging.info("No YARA rules were triggered. Nothing will be written to the output archive.") - -class _TarLz4Wrapper: - - def __init__(self, path) -> None: - self._lz4 = lz4.frame.open(path, 'wb') - self._tar = tarfile.open(fileobj=self._lz4f, mode="w") - - def writestr(self, path: str, value: str | bytes): - info = tarfile.TarInfo(path) - info.size = len(value) - self._tar.addfile(info, io.BytesIO(value)) - - def write(self, path: str, arcname: str): - self._tar.add(path, arcname) - - def __enter__(self): - return self - - def __exit__(self, type, value, traceback): - self._tar.close() - self._lz4.close() \ No newline at end of file From e3c062d987965587d00d341f501c3a62e1c6ae05 Mon Sep 17 00:00:00 2001 From: John Lugton Date: Fri, 20 Sep 2024 15:55:08 +0100 Subject: [PATCH 5/6] Fix build failures --- tests/test_import.py | 5 ++--- tests/test_linux/test_base.py | 9 ++------- tests/test_shared.py | 5 ++--- varc_core/systems/base_system.py | 10 +++++----- 4 files changed, 11 insertions(+), 18 deletions(-) diff --git a/tests/test_import.py b/tests/test_import.py index b7c3062..ef442ec 100644 --- a/tests/test_import.py +++ b/tests/test_import.py @@ -8,12 +8,11 @@ class TestImport(unittest.TestCase): system: BaseSystem - zip_path: str @classmethod def setUpClass(cls) -> None: cls.system = acquire_system() - cls.zip_path = cls.system.acquire_volatile() + cls.system.acquire_volatile() @classmethod def tearDownClass(cls) -> None: @@ -29,5 +28,5 @@ def test_get_network(self) -> None: def test_got_files(self) -> None: # Check we got atleast 10 files - with ZipFile(self.zip_path) as z: + with ZipFile(self.system.output_path) as z: self.assertGreater(len(z.namelist()), 10) diff --git a/tests/test_linux/test_base.py b/tests/test_linux/test_base.py index 1d26c06..0cd73fd 100644 --- a/tests/test_linux/test_base.py +++ b/tests/test_linux/test_base.py @@ -5,18 +5,13 @@ class TestBaseCases(unittest.TestCase): system: BaseSystem - zip_path: str @classmethod def setUpClass(cls) -> None: cls.system = acquire_system() - cls.zip_path = cls.system.acquire_volatile() + cls.system.acquire_volatile() - @classmethod - def tearDownClass(cls) -> None: - pass - def test_some_processes(self) -> None: processes = self.system.get_processes() self.assertTrue(len(processes) > 0) @@ -28,7 +23,7 @@ def test_dump_files(self) -> None: open_files = self.system.dump_loaded_files() self.assertTrue(len(open_files) > 0) # Check we pulled at least one file from /bin/ - with ZipFile(self.zip_path) as z: + with ZipFile(self.system.output_path) as z: binary_files = [binary for binary in z.namelist() if ("/bin/" in binary in binary.lower())] self.assertGreater(len(binary_files), 0) \ No newline at end of file diff --git a/tests/test_shared.py b/tests/test_shared.py index 252582b..a4781a6 100644 --- a/tests/test_shared.py +++ b/tests/test_shared.py @@ -6,12 +6,11 @@ class TestBaseCases(unittest.TestCase): system: BaseSystem - zip_path: str @classmethod def setUpClass(cls) -> None: cls.system = acquire_system() - cls.zip_path = cls.system.acquire_volatile() + cls.system.acquire_volatile() @classmethod def tearDownClass(cls) -> None: @@ -27,5 +26,5 @@ def test_get_network(self) -> None: def test_got_files(self) -> None: # Check we got atleast 10 files - with ZipFile(self.zip_path) as z: + with ZipFile(self.system.output_path) as z: self.assertGreater(len(z.namelist()), 10) diff --git a/varc_core/systems/base_system.py b/varc_core/systems/base_system.py index 722bf51..cd9a6c5 100644 --- a/varc_core/systems/base_system.py +++ b/varc_core/systems/base_system.py @@ -18,7 +18,7 @@ import zipfile from base64 import b64encode from datetime import datetime -from typing import Any, List, Optional +from typing import IO, Any, List, Optional, Union import lz4.frame import mss @@ -40,12 +40,12 @@ class _TarLz4Wrapper: def __init__(self, path) -> None: self._lz4 = lz4.frame.open(path, 'wb') - self._tar = tarfile.open(fileobj=self._lz4f, mode="w") + self._tar = tarfile.open(fileobj=self._lz4, mode="w") - def writestr(self, path: str, value: str | bytes): + def writestr(self, path: str, value: Union[str, bytes]): info = tarfile.TarInfo(path) info.size = len(value) - self._tar.addfile(info, io.BytesIO(value)) + self._tar.addfile(info, io.BytesIO(value if isinstance(value, bytes) else value.encode())) def write(self, path: str, arcname: str): self._tar.add(path, arcname) @@ -339,7 +339,7 @@ def acquire_volatile(self) -> None: except FileNotFoundError: logging.warning(f"Could not open {file_path} for reading") - def _open_output(self) -> zipfile.ZipFile | _TarLz4Wrapper: + def _open_output(self) -> Union[zipfile.ZipFile, _TarLz4Wrapper]: if self.output_path.endswith('.tar.lz4'): return _TarLz4Wrapper(self.output_path) else: From 2122d03339757892675c8aa11d2d24c0287d563f Mon Sep 17 00:00:00 2001 From: John Lugton Date: Fri, 20 Sep 2024 16:05:38 +0100 Subject: [PATCH 6/6] Fix types --- varc_core/systems/base_system.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/varc_core/systems/base_system.py b/varc_core/systems/base_system.py index cd9a6c5..a9b7a92 100644 --- a/varc_core/systems/base_system.py +++ b/varc_core/systems/base_system.py @@ -18,9 +18,9 @@ import zipfile from base64 import b64encode from datetime import datetime -from typing import IO, Any, List, Optional, Union +from typing import Any, List, Optional, Union -import lz4.frame +import lz4.frame # type: ignore import mss import psutil from tqdm import tqdm @@ -38,22 +38,22 @@ class _TarLz4Wrapper: - def __init__(self, path) -> None: + def __init__(self, path: str) -> None: self._lz4 = lz4.frame.open(path, 'wb') self._tar = tarfile.open(fileobj=self._lz4, mode="w") - def writestr(self, path: str, value: Union[str, bytes]): + def writestr(self, path: str, value: Union[str, bytes]) -> None: info = tarfile.TarInfo(path) info.size = len(value) self._tar.addfile(info, io.BytesIO(value if isinstance(value, bytes) else value.encode())) - def write(self, path: str, arcname: str): + def write(self, path: str, arcname: str) -> None: self._tar.add(path, arcname) - def __enter__(self): + def __enter__(self) -> "_TarLz4Wrapper": return self - def __exit__(self, type, value, traceback): + def __exit__(self, type: Any, value: Any, traceback: Any) -> None: self._tar.close() self._lz4.close()