diff --git a/disassemblers/ofrak_binary_ninja/ofrak_binary_ninja/components/binary_ninja_analyzer.py b/disassemblers/ofrak_binary_ninja/ofrak_binary_ninja/components/binary_ninja_analyzer.py index 0ca5a0234..b9d959dd7 100644 --- a/disassemblers/ofrak_binary_ninja/ofrak_binary_ninja/components/binary_ninja_analyzer.py +++ b/disassemblers/ofrak_binary_ninja/ofrak_binary_ninja/components/binary_ninja_analyzer.py @@ -1,5 +1,4 @@ import logging -import tempfile from dataclasses import dataclass from typing import Optional, List @@ -29,11 +28,9 @@ async def analyze( self, resource: Resource, config: Optional[BinaryNinjaAnalyzerConfig] = None ) -> BinaryNinjaAnalysis: if not config: - resource_data = await resource.get_data() - temp_file = tempfile.NamedTemporaryFile() - temp_file.write(resource_data) - temp_file.flush() - bv = open_view(temp_file.name) + async with resource.temp_to_disk(delete=False) as temp_path: + bv = open_view(temp_path) + return BinaryNinjaAnalysis(bv) else: bv = BinaryViewType.get_view_of_file(config.bndb_file) diff --git a/ofrak_core/CHANGELOG.md b/ofrak_core/CHANGELOG.md index f728379f9..4f68e2897 100644 --- a/ofrak_core/CHANGELOG.md +++ b/ofrak_core/CHANGELOG.md @@ -35,6 +35,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) - Use PyPI version of `bincopy`, upgrade to version 20.0.0 ([#528](https://github.com/redballoonsecurity/ofrak/pull/528)) - Fix bugs on Windows arising from using `os.path` methods when only forward-slashes are acceptable ([#521](https://github.com/redballoonsecurity/ofrak/pull/521)) - Made some changes to OFRAK test suite to improve test coverage on Windows ([#487](https://github.com/redballoonsecurity/ofrak/pull/487)) +- Fix usage of `NamedTemporaryFile` with external tools on Windows ([#486](https://github.com/redballoonsecurity/ofrak/pull/486)) ### Changed - By default, the ofrak log is now `ofrak-YYYYMMDDhhmmss.log` rather than just `ofrak.log` and the name can be specified on the command line ([#480](https://github.com/redballoonsecurity/ofrak/pull/480)) diff --git a/ofrak_core/ofrak/core/apk.py b/ofrak_core/ofrak/core/apk.py index adb4689ff..949fa1133 100644 --- a/ofrak_core/ofrak/core/apk.py +++ b/ofrak_core/ofrak/core/apk.py @@ -2,7 +2,7 @@ import os import pathlib import sys -import tempfile +import tempfile312 as tempfile from subprocess import CalledProcessError from dataclasses import dataclass @@ -101,10 +101,7 @@ async def unpack(self, resource: Resource, config=None): :param config: """ apk = await resource.view_as(Apk) - data = await resource.get_data() - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(data) - temp_file.flush() + async with resource.temp_to_disk() as temp_path: with tempfile.TemporaryDirectory() as temp_flush_dir: cmd = [ "apktool", @@ -112,7 +109,7 @@ async def unpack(self, resource: Resource, config=None): "--output", temp_flush_dir, "--force", - temp_file.name, + temp_path, ] proc = await asyncio.create_subprocess_exec( *cmd, @@ -156,7 +153,8 @@ async def pack( apk = await resource.view_as(Apk) temp_flush_dir = await apk.flush_to_disk() apk_suffix = ".apk" - with tempfile.NamedTemporaryFile(suffix=apk_suffix) as temp_apk: + with tempfile.NamedTemporaryFile(suffix=apk_suffix, delete_on_close=False) as temp_apk: + temp_apk.close() apk_cmd = [ "apktool", "build", @@ -218,13 +216,11 @@ async def identify(self, resource: Resource, config=None) -> None: if magic.mime == "application/vnd.android.package-archive": resource.add_tag(Apk) elif magic is not None and magic.mime in ["application/java-archive", "application/zip"]: - with tempfile.NamedTemporaryFile(suffix=".zip") as temp_file: - temp_file.write(await resource.get_data()) - temp_file.flush() + async with resource.temp_to_disk(suffix=".zip") as temp_path: unzip_cmd = [ "unzip", "-l", - temp_file.name, + temp_path, ] unzip_proc = await asyncio.create_subprocess_exec( *unzip_cmd, diff --git a/ofrak_core/ofrak/core/binwalk.py b/ofrak_core/ofrak/core/binwalk.py index 6a87d3723..e433d4e8e 100644 --- a/ofrak_core/ofrak/core/binwalk.py +++ b/ofrak_core/ofrak/core/binwalk.py @@ -1,5 +1,4 @@ import asyncio -import tempfile from concurrent.futures.process import ProcessPoolExecutor from dataclasses import dataclass from typing import Dict @@ -62,15 +61,11 @@ def __init__( async def analyze(self, resource: Resource, config=None) -> BinwalkAttributes: if not BINWALK_INSTALLED: raise ComponentMissingDependencyError(self, BINWALK_TOOL) - with tempfile.NamedTemporaryFile() as temp_file: - data = await resource.get_data() - temp_file.write(data) - temp_file.flush() - + async with resource.temp_to_disk() as temp_path: # Should errors be handled the way they are in the `DataSummaryAnalyzer`? Likely to be # overkill here. offsets = await asyncio.get_running_loop().run_in_executor( - self.pool, _run_binwalk_on_file, temp_file.name + self.pool, _run_binwalk_on_file, temp_path ) return BinwalkAttributes(offsets) diff --git a/ofrak_core/ofrak/core/cpio.py b/ofrak_core/ofrak/core/cpio.py index 4557efd29..685701990 100644 --- a/ofrak_core/ofrak/core/cpio.py +++ b/ofrak_core/ofrak/core/cpio.py @@ -1,6 +1,6 @@ import asyncio import logging -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from enum import Enum from subprocess import CalledProcessError @@ -100,8 +100,8 @@ async def unpack(self, resource: Resource, config=None): cwd=temp_flush_dir, ) await proc.communicate(input=resource_data) - if proc.returncode: - raise CalledProcessError(returncode=proc.returncode, cmd=cmd) + # if proc.returncode: + # raise CalledProcessError(returncode=proc.returncode, cmd=cmd) await cpio_v.initialize_from_disk(temp_flush_dir) diff --git a/ofrak_core/ofrak/core/elf/lief_modifier.py b/ofrak_core/ofrak/core/elf/lief_modifier.py index a6774eed3..04633ce8d 100644 --- a/ofrak_core/ofrak/core/elf/lief_modifier.py +++ b/ofrak_core/ofrak/core/elf/lief_modifier.py @@ -1,9 +1,9 @@ -import tempfile from dataclasses import dataclass from typing import List, Optional import lief +import tempfile312 as tempfile from ofrak.component.modifier import Modifier from ofrak.model.component_model import ComponentConfig from ofrak.resource import Resource @@ -66,9 +66,9 @@ async def modify(self, resource: Resource, config: LiefAddSegmentConfig) -> None else: _ = binary.add(segment) - with tempfile.NamedTemporaryFile() as temp_file: + with tempfile.NamedTemporaryFile(delete_on_close=False) as temp_file: + temp_file.close() binary.write(temp_file.name) - temp_file.flush() with open(temp_file.name, "rb") as f_handle: new_data = f_handle.read() # replace all old content (old range) with new content from Lief @@ -93,9 +93,9 @@ async def modify(self, resource: Resource, config: LiefAddSectionModifierConfig) section.flags = config.flags binary.add(section) - with tempfile.NamedTemporaryFile() as temp_file: + with tempfile.NamedTemporaryFile(delete_on_close=False) as temp_file: + temp_file.close() binary.write(temp_file.name) - temp_file.flush() with open(temp_file.name, "rb") as f_handle: new_data = f_handle.read() # replace all old content (old range) with new content from Lief @@ -117,9 +117,9 @@ async def modify(self, resource: Resource, config: LiefRemoveSectionModifierConf raise AttributeError(f"No section with name {config.name}") binary.remove(section) - with tempfile.NamedTemporaryFile() as temp_file: + with tempfile.NamedTemporaryFile(delete_on_close=False) as temp_file: + temp_file.close() binary.write(temp_file.name) - temp_file.flush() with open(temp_file.name, "rb") as f_handle: new_data = f_handle.read() # replace all old content (old range) with new content from Lief diff --git a/ofrak_core/ofrak/core/extfs.py b/ofrak_core/ofrak/core/extfs.py index b09cd4001..3c50705b9 100644 --- a/ofrak_core/ofrak/core/extfs.py +++ b/ofrak_core/ofrak/core/extfs.py @@ -1,5 +1,5 @@ import asyncio -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError @@ -55,16 +55,13 @@ class ExtUnpacker(Unpacker[None]): external_dependencies = (_DEBUGFS,) async def unpack(self, resource: Resource, config: ComponentConfig = None) -> None: - with tempfile.NamedTemporaryFile(suffix=".extfs") as temp_fs_file: - temp_fs_file.write(await resource.get_data()) - temp_fs_file.flush() - + async with resource.temp_to_disk(suffix=".extfs") as temp_fs_path: with tempfile.TemporaryDirectory() as temp_dir: command = [ "debugfs", "-R", f"rdump / {temp_dir}", - temp_fs_file.name, + temp_fs_path, ] proc = await asyncio.create_subprocess_exec( *command, diff --git a/ofrak_core/ofrak/core/filesystem.py b/ofrak_core/ofrak/core/filesystem.py index a6849b609..d62c8836f 100644 --- a/ofrak_core/ofrak/core/filesystem.py +++ b/ofrak_core/ofrak/core/filesystem.py @@ -2,7 +2,7 @@ import stat import sys import warnings -import tempfile +import tempfile312 as tempfile import posixpath from pathlib import PurePath from dataclasses import dataclass diff --git a/ofrak_core/ofrak/core/gzip.py b/ofrak_core/ofrak/core/gzip.py index 342810a6c..66df3d5ee 100644 --- a/ofrak_core/ofrak/core/gzip.py +++ b/ofrak_core/ofrak/core/gzip.py @@ -3,7 +3,7 @@ from typing import Optional import zlib from subprocess import CalledProcessError -import tempfile +import tempfile312 as tempfile from ofrak.component.packer import Packer from ofrak.component.unpacker import Unpacker @@ -110,9 +110,9 @@ async def pack_with_zlib_module(data: bytes) -> bytes: @staticmethod async def pack_with_pigz(data: bytes) -> bytes: - with tempfile.NamedTemporaryFile() as uncompressed_file: + with tempfile.NamedTemporaryFile(delete_on_close=False) as uncompressed_file: uncompressed_file.write(data) - uncompressed_file.flush() + uncompressed_file.close() cmd = [ "pigz", diff --git a/ofrak_core/ofrak/core/iso9660.py b/ofrak_core/ofrak/core/iso9660.py index dcbb05113..4549846b3 100644 --- a/ofrak_core/ofrak/core/iso9660.py +++ b/ofrak_core/ofrak/core/iso9660.py @@ -1,7 +1,7 @@ import asyncio import logging +import tempfile312 as tempfile import posixpath -import tempfile from dataclasses import dataclass from io import BytesIO from subprocess import CalledProcessError @@ -301,7 +301,8 @@ async def pack(self, resource: Resource, config=None) -> None: iso_attrs = resource.get_attributes(ISO9660ImageAttributes) temp_flush_dir = await iso_view.flush_to_disk() - with tempfile.NamedTemporaryFile(suffix=".iso", mode="rb") as temp: + with tempfile.NamedTemporaryFile(suffix=".iso", mode="rb", delete_on_close=False) as temp: + temp.close() cmd = [ "mkisofs", *(["-J"] if iso_attrs.has_joliet else []), @@ -329,7 +330,8 @@ async def pack(self, resource: Resource, config=None) -> None: returncode = await proc.wait() if proc.returncode: raise CalledProcessError(returncode=returncode, cmd=cmd) - new_data = temp.read() + with open(temp.name, "rb") as new_fh: + new_data = new_fh.read() # Passing in the original range effectively replaces the original data with the new data resource.queue_patch(Range(0, await resource.get_data_length()), new_data) diff --git a/ofrak_core/ofrak/core/jffs2.py b/ofrak_core/ofrak/core/jffs2.py index 48883a095..1b0b94c1a 100644 --- a/ofrak_core/ofrak/core/jffs2.py +++ b/ofrak_core/ofrak/core/jffs2.py @@ -1,6 +1,6 @@ import asyncio import logging -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError @@ -38,18 +38,14 @@ class Jffs2Unpacker(Unpacker[None]): external_dependencies = (JEFFERSON,) async def unpack(self, resource: Resource, config=None): - with tempfile.NamedTemporaryFile() as temp_file: - resource_data = await resource.get_data() - temp_file.write(resource_data) - temp_file.flush() - + async with resource.temp_to_disk() as temp_path: with tempfile.TemporaryDirectory() as temp_flush_dir: cmd = [ "jefferson", "--force", "--dest", temp_flush_dir, - temp_file.name, + temp_path, ] proc = await asyncio.create_subprocess_exec( *cmd, @@ -73,7 +69,8 @@ class Jffs2Packer(Packer[None]): async def pack(self, resource: Resource, config=None): jffs2_view: Jffs2Filesystem = await resource.view_as(Jffs2Filesystem) temp_flush_dir = await jffs2_view.flush_to_disk() - with tempfile.NamedTemporaryFile(suffix=".sqsh", mode="rb") as temp: + with tempfile.NamedTemporaryFile(suffix=".sqsh", mode="rb", delete_on_close=False) as temp: + temp.close() cmd = [ "mkfs.jffs2", "-r", @@ -87,7 +84,8 @@ async def pack(self, resource: Resource, config=None): returncode = await proc.wait() if proc.returncode: raise CalledProcessError(returncode=returncode, cmd=cmd) - new_data = temp.read() + with open(temp.name, "rb") as new_fh: + new_data = new_fh.read() # Passing in the original range effectively replaces the original data with the new data resource.queue_patch(Range(0, await resource.get_data_length()), new_data) diff --git a/ofrak_core/ofrak/core/lzo.py b/ofrak_core/ofrak/core/lzo.py index f69634c52..ac22da72c 100644 --- a/ofrak_core/ofrak/core/lzo.py +++ b/ofrak_core/ofrak/core/lzo.py @@ -1,5 +1,4 @@ import asyncio -import tempfile from subprocess import CalledProcessError from ofrak.component.packer import Packer @@ -37,27 +36,18 @@ class LzoUnpacker(Unpacker[None]): external_dependencies = (LZOP,) async def unpack(self, resource: Resource, config: ComponentConfig = None) -> None: - with tempfile.NamedTemporaryFile(suffix=".lzo") as compressed_file: - compressed_file.write(await resource.get_data()) - compressed_file.flush() - - cmd = [ - "lzop", - "-d", - "-f", - "-c", - compressed_file.name, - ] - proc = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await proc.communicate() - if proc.returncode: - raise CalledProcessError(returncode=proc.returncode, cmd=cmd) - - await resource.create_child(tags=(GenericBinary,), data=stdout) + cmd = ["lzop", "-d", "-f"] + proc = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate(await resource.get_data()) + if proc.returncode: + raise CalledProcessError(returncode=proc.returncode, cmd=cmd) + + await resource.create_child(tags=(GenericBinary,), data=stdout) class LzoPacker(Packer[None]): @@ -73,28 +63,20 @@ async def pack(self, resource: Resource, config: ComponentConfig = None): child_file = await lzo_view.get_child() uncompressed_data = await child_file.resource.get_data() - with tempfile.NamedTemporaryFile(suffix=".lzo") as uncompressed_file: - uncompressed_file.write(uncompressed_data) - uncompressed_file.flush() - - cmd = [ - "lzop", - "-f", - "-c", - uncompressed_file.name, - ] - proc = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await proc.communicate() - if proc.returncode: - raise CalledProcessError(returncode=proc.returncode, cmd=cmd) - - compressed_data = stdout - original_size = await lzo_view.resource.get_data_length() - resource.queue_patch(Range(0, original_size), compressed_data) + cmd = ["lzop", "-f"] + proc = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate(uncompressed_data) + if proc.returncode: + raise CalledProcessError(returncode=proc.returncode, cmd=cmd) + + compressed_data = stdout + original_size = await lzo_view.resource.get_data_length() + resource.queue_patch(Range(0, original_size), compressed_data) MagicMimeIdentifier.register(LzoData, "application/x-lzop") diff --git a/ofrak_core/ofrak/core/patch_maker/modifiers.py b/ofrak_core/ofrak/core/patch_maker/modifiers.py index 76f1c9171..a5c0c098f 100644 --- a/ofrak_core/ofrak/core/patch_maker/modifiers.py +++ b/ofrak_core/ofrak/core/patch_maker/modifiers.py @@ -1,7 +1,7 @@ import asyncio import logging import os -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from typing import Dict, List, Optional, Tuple, Type, Union, cast diff --git a/ofrak_core/ofrak/core/rar.py b/ofrak_core/ofrak/core/rar.py index 6def88610..09eede84a 100644 --- a/ofrak_core/ofrak/core/rar.py +++ b/ofrak_core/ofrak/core/rar.py @@ -1,5 +1,5 @@ import asyncio -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError @@ -30,7 +30,7 @@ class RarArchive(GenericBinary, FilesystemRoot): class RarUnpacker(Unpacker[None]): """ - Unpack RAR archives using the free `unrar` tool. + Unpack RAR archives using the free `unar` tool. """ targets = (RarArchive,) @@ -38,28 +38,24 @@ class RarUnpacker(Unpacker[None]): external_dependencies = (UNAR,) async def unpack(self, resource: Resource, config: ComponentConfig = None): - with tempfile.NamedTemporaryFile( - suffix=".rar" - ) as temp_archive, tempfile.TemporaryDirectory() as temp_dir: - temp_archive.write(await resource.get_data()) - temp_archive.flush() - - cmd = [ - "unar", - "-no-directory", - "-no-recursion", - temp_archive.name, - ] - proc = await asyncio.create_subprocess_exec( - *cmd, - cwd=temp_dir, - ) - returncode = await proc.wait() - if proc.returncode: - raise CalledProcessError(returncode=returncode, cmd=cmd) - - rar_view = await resource.view_as(RarArchive) - await rar_view.initialize_from_disk(temp_dir) + async with resource.temp_to_disk(suffix=".rar") as temp_archive: + with tempfile.TemporaryDirectory() as temp_dir: + cmd = [ + "unar", + "-no-directory", + "-no-recursion", + temp_archive, + ] + proc = await asyncio.create_subprocess_exec( + *cmd, + cwd=temp_dir, + ) + returncode = await proc.wait() + if proc.returncode: + raise CalledProcessError(returncode=returncode, cmd=cmd) + + rar_view = await resource.view_as(RarArchive) + await rar_view.initialize_from_disk(temp_dir) MagicMimeIdentifier.register(RarArchive, "application/x-rar-compressed") diff --git a/ofrak_core/ofrak/core/seven_zip.py b/ofrak_core/ofrak/core/seven_zip.py index 64d5b7688..edd2d071a 100644 --- a/ofrak_core/ofrak/core/seven_zip.py +++ b/ofrak_core/ofrak/core/seven_zip.py @@ -1,7 +1,7 @@ import asyncio import logging import os -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError @@ -39,15 +39,13 @@ class SevenZUnpacker(Unpacker[None]): async def unpack(self, resource: Resource, config=None): seven_zip_v = await resource.view_as(SevenZFilesystem) resource_data = await seven_zip_v.resource.get_data() - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(resource_data) - temp_file.flush() + async with resource.temp_to_disk(suffix=".7z") as temp_path: with tempfile.TemporaryDirectory() as temp_flush_dir: cmd = [ "7zz", "x", f"-o{temp_flush_dir}", - temp_file.name, + temp_path, ] proc = await asyncio.create_subprocess_exec( *cmd, diff --git a/ofrak_core/ofrak/core/squashfs.py b/ofrak_core/ofrak/core/squashfs.py index 10bdf645a..3f73c2c41 100644 --- a/ofrak_core/ofrak/core/squashfs.py +++ b/ofrak_core/ofrak/core/squashfs.py @@ -1,6 +1,6 @@ import asyncio import logging -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError @@ -66,11 +66,7 @@ class SquashfsUnpacker(Unpacker[None]): external_dependencies = (UNSQUASHFS,) async def unpack(self, resource: Resource, config=None): - with tempfile.NamedTemporaryFile() as temp_file: - resource_data = await resource.get_data() - temp_file.write(resource_data) - temp_file.flush() - + async with resource.temp_to_disk() as temp_path: with tempfile.TemporaryDirectory() as temp_flush_dir: cmd = [ "unsquashfs", @@ -78,7 +74,7 @@ async def unpack(self, resource: Resource, config=None): "-force", "-dest", temp_flush_dir, - temp_file.name, + temp_path, ] proc = await asyncio.create_subprocess_exec( *cmd, @@ -102,7 +98,8 @@ class SquashfsPacker(Packer[None]): async def pack(self, resource: Resource, config=None): squashfs_view: SquashfsFilesystem = await resource.view_as(SquashfsFilesystem) temp_flush_dir = await squashfs_view.flush_to_disk() - with tempfile.NamedTemporaryFile(suffix=".sqsh", mode="rb") as temp: + with tempfile.NamedTemporaryFile(suffix=".sqsh", mode="rb", delete_on_close=False) as temp: + temp.close() cmd = [ "mksquashfs", temp_flush_dir, @@ -115,7 +112,8 @@ async def pack(self, resource: Resource, config=None): returncode = await proc.wait() if proc.returncode: raise CalledProcessError(returncode=returncode, cmd=cmd) - new_data = temp.read() + with open(temp.name, "rb") as new_fh: + new_data = new_fh.read() # Passing in the original range effectively replaces the original data with the new data resource.queue_patch(Range(0, await resource.get_data_length()), new_data) diff --git a/ofrak_core/ofrak/core/strings_analyzer.py b/ofrak_core/ofrak/core/strings_analyzer.py index cd0b40243..4bc7fe9e6 100644 --- a/ofrak_core/ofrak/core/strings_analyzer.py +++ b/ofrak_core/ofrak/core/strings_analyzer.py @@ -1,5 +1,4 @@ import asyncio -import tempfile from dataclasses import dataclass from typing import Dict, Optional @@ -61,16 +60,13 @@ async def analyze( config = StringsAnalyzerConfig() strings = dict() - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(await resource.get_data()) - temp_file.flush() - + async with resource.temp_to_disk() as temp_path: proc = await asyncio.subprocess.create_subprocess_exec( "strings", "-t", "d", f"-{config.min_length}", - temp_file.name, + temp_path, stdout=asyncio.subprocess.PIPE, ) diff --git a/ofrak_core/ofrak/core/tar.py b/ofrak_core/ofrak/core/tar.py index fabe7f876..6c023323a 100644 --- a/ofrak_core/ofrak/core/tar.py +++ b/ofrak_core/ofrak/core/tar.py @@ -1,6 +1,6 @@ import asyncio import os.path -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError @@ -37,16 +37,13 @@ class TarUnpacker(Unpacker[None]): async def unpack(self, resource: Resource, config: ComponentConfig = None) -> None: # Write the archive data to a file - with tempfile.NamedTemporaryFile(suffix=".tar") as temp_archive: - temp_archive.write(await resource.get_data()) - temp_archive.flush() - + async with resource.temp_to_disk(suffix=".tar") as temp_archive_path: # Check the archive member files to ensure none unpack to a parent directory cmd = [ "tar", "-P", "-tf", - temp_archive.name, + temp_archive_path, ] proc = await asyncio.create_subprocess_exec( *cmd, @@ -67,7 +64,7 @@ async def unpack(self, resource: Resource, config: ComponentConfig = None) -> No # Unpack into a temporary directory using the temporary file with tempfile.TemporaryDirectory() as temp_dir: - command = ["tar", "--xattrs", "-C", temp_dir, "-xf", temp_archive.name] + command = ["tar", "--xattrs", "-C", temp_dir, "-xf", temp_archive_path] proc = await asyncio.create_subprocess_exec( *command, ) @@ -94,7 +91,8 @@ async def pack(self, resource: Resource, config: ComponentConfig = None) -> None flush_dir = await tar_view.flush_to_disk() # Pack it back into a temporary archive - with tempfile.NamedTemporaryFile(suffix=".tar") as temp_archive: + with tempfile.NamedTemporaryFile(suffix=".tar", delete_on_close=False) as temp_archive: + temp_archive.close() cmd = [ "tar", "--xattrs", @@ -112,7 +110,8 @@ async def pack(self, resource: Resource, config: ComponentConfig = None) -> None raise CalledProcessError(returncode=returncode, cmd=cmd) # Replace the original archive data - resource.queue_patch(Range(0, await resource.get_data_length()), temp_archive.read()) + with open(temp_archive.name, "rb") as new_fh: + resource.queue_patch(Range(0, await resource.get_data_length()), new_fh.read()) MagicMimeIdentifier.register(TarArchive, "application/x-tar") diff --git a/ofrak_core/ofrak/core/ubi.py b/ofrak_core/ofrak/core/ubi.py index f82a64cda..f59dc72b7 100644 --- a/ofrak_core/ofrak/core/ubi.py +++ b/ofrak_core/ofrak/core/ubi.py @@ -1,5 +1,5 @@ import asyncio -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass import logging from typing import List, Tuple @@ -131,15 +131,11 @@ class UbiAnalyzer(Analyzer[None, Ubi]): async def analyze(self, resource: Resource, config=None) -> Ubi: # Flush to disk - with tempfile.NamedTemporaryFile() as temp_file: - resource_data = await resource.get_data() - temp_file.write(resource_data) - temp_file.flush() - + async with resource.temp_to_disk() as temp_path: ubi_obj = ubireader_ubi( ubi_io.ubi_file( - temp_file.name, - block_size=guess_peb_size(temp_file.name), + temp_path, + block_size=guess_peb_size(temp_path), start_offset=0, end_offset=None, ) @@ -241,7 +237,6 @@ class UbiPacker(Packer[None]): async def pack(self, resource: Resource, config=None) -> None: ubi_view = await resource.view_as(Ubi) - # with tempfile.NamedTemporaryFile(mode="rb") as temp: with tempfile.TemporaryDirectory() as temp_flush_dir: ubi_volumes = await resource.get_children() ubinize_ini_entries = [] diff --git a/ofrak_core/ofrak/core/ubifs.py b/ofrak_core/ofrak/core/ubifs.py index 9af8157ff..250055125 100644 --- a/ofrak_core/ofrak/core/ubifs.py +++ b/ofrak_core/ofrak/core/ubifs.py @@ -1,5 +1,5 @@ import asyncio -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass import logging from subprocess import CalledProcessError @@ -97,15 +97,11 @@ class UbifsAnalyzer(Analyzer[None, Ubifs]): external_dependencies = (PY_LZO_TOOL,) async def analyze(self, resource: Resource, config=None) -> Ubifs: - with tempfile.NamedTemporaryFile() as temp_file: - resource_data = await resource.get_data() - temp_file.write(resource_data) - temp_file.flush() - + async with resource.temp_to_disk() as temp_path: ubifs_obj = ubireader_ubifs( ubi_io.ubi_file( - temp_file.name, - block_size=guess_leb_size(temp_file.name), + temp_path, + block_size=guess_leb_size(temp_path), start_offset=0, end_offset=None, ) @@ -172,7 +168,8 @@ async def pack(self, resource: Resource, config=None) -> None: ubifs_view = await resource.view_as(Ubifs) flush_dir = await ubifs_view.flush_to_disk() - with tempfile.NamedTemporaryFile(mode="rb") as temp: + with tempfile.NamedTemporaryFile(mode="rb", delete_on_close=False) as temp: + temp.close() cmd = [ "mkfs.ubifs", "-m", @@ -202,7 +199,8 @@ async def pack(self, resource: Resource, config=None) -> None: returncode = await proc.wait() if proc.returncode: raise CalledProcessError(returncode=returncode, cmd=cmd) - new_data = temp.read() + with open(temp.name, "rb") as new_fh: + new_data = new_fh.read() resource.queue_patch(Range(0, await resource.get_data_length()), new_data) diff --git a/ofrak_core/ofrak/core/uefi.py b/ofrak_core/ofrak/core/uefi.py index 887b8dd18..7186127be 100644 --- a/ofrak_core/ofrak/core/uefi.py +++ b/ofrak_core/ofrak/core/uefi.py @@ -1,7 +1,7 @@ import os import asyncio import logging -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError diff --git a/ofrak_core/ofrak/core/xattr_stub.py b/ofrak_core/ofrak/core/xattr_stub.py index 2015bddd0..b6a27ccac 100644 --- a/ofrak_core/ofrak/core/xattr_stub.py +++ b/ofrak_core/ofrak/core/xattr_stub.py @@ -141,7 +141,7 @@ def removexattr(f, attr, symlink=False): def _warn_user_no_xattr(function_name: str) -> None: LOGGER.warning( - f"Function {function_name} not found. Library xattr is not available on Windows platforms. \ - Extended attributes will not be properly handled while using OFRAK on this platform. \ - If you require extended attributes, please use a platform that supports xattr." + f"Function {function_name} not found. Library xattr is not available on Windows platforms. " + "Extended attributes will not be properly handled while using OFRAK on this platform. " + "If you require extended attributes, please use a platform that supports xattr." ) diff --git a/ofrak_core/ofrak/core/zip.py b/ofrak_core/ofrak/core/zip.py index 2a19d06a8..39419a185 100644 --- a/ofrak_core/ofrak/core/zip.py +++ b/ofrak_core/ofrak/core/zip.py @@ -1,7 +1,7 @@ import asyncio import logging import os -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from subprocess import CalledProcessError @@ -51,13 +51,11 @@ class ZipUnpacker(Unpacker[None]): async def unpack(self, resource: Resource, config=None): zip_view = await resource.view_as(ZipArchive) - with tempfile.NamedTemporaryFile(suffix=".zip") as temp_archive: - temp_archive.write(await resource.get_data()) - temp_archive.flush() + async with resource.temp_to_disk(suffix=".zip") as temp_path: with tempfile.TemporaryDirectory() as temp_dir: cmd = [ "unzip", - temp_archive.name, + temp_path, "-d", temp_dir, ] diff --git a/ofrak_core/ofrak/core/zstd.py b/ofrak_core/ofrak/core/zstd.py index f0cddee81..91df1c1b9 100644 --- a/ofrak_core/ofrak/core/zstd.py +++ b/ofrak_core/ofrak/core/zstd.py @@ -1,5 +1,4 @@ import asyncio -import tempfile from dataclasses import dataclass from typing import Optional from subprocess import CalledProcessError @@ -42,29 +41,15 @@ class ZstdUnpacker(Unpacker[None]): external_dependencies = (ZSTD,) async def unpack(self, resource: Resource, config: ComponentConfig = None) -> None: - with tempfile.NamedTemporaryFile(suffix=".zstd") as compressed_file: - compressed_file.write(await resource.get_data()) - compressed_file.flush() - output_filename = tempfile.mktemp() - - cmd = [ - "zstd", - "-d", - "-k", - compressed_file.name, - "-o", - output_filename, - ] - proc = await asyncio.create_subprocess_exec( - *cmd, - ) - returncode = await proc.wait() - if proc.returncode: - raise CalledProcessError(returncode=returncode, cmd=cmd) - with open(output_filename, "rb") as f: - result = f.read() - - await resource.create_child(tags=(GenericBinary,), data=result) + cmd = ["zstd", "-d", "-k"] + proc = await asyncio.create_subprocess_exec( + *cmd, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE + ) + result, _ = await proc.communicate(await resource.get_data()) + if proc.returncode: + raise CalledProcessError(returncode=proc.returncode, cmd=cmd) + + await resource.create_child(tags=(GenericBinary,), data=result) class ZstdPacker(Packer[ZstdPackerConfig]): @@ -82,27 +67,19 @@ async def pack(self, resource: Resource, config: Optional[ZstdPackerConfig] = No child_file = await zstd_view.get_child() uncompressed_data = await child_file.resource.get_data() - with tempfile.NamedTemporaryFile() as uncompressed_file: - uncompressed_file.write(uncompressed_data) - uncompressed_file.flush() - output_filename = tempfile.mktemp() - - command = ["zstd", "-T0", f"-{config.compression_level}"] - if config.compression_level > 19: - command.append("--ultra") - command.extend([uncompressed_file.name, "-o", output_filename]) - proc = await asyncio.create_subprocess_exec( - *command, - ) - returncode = await proc.wait() - if proc.returncode: - raise CalledProcessError(returncode=returncode, cmd=command) - with open(output_filename, "rb") as f: - result = f.read() - - compressed_data = result - original_size = await zstd_view.resource.get_data_length() - resource.queue_patch(Range(0, original_size), compressed_data) + command = ["zstd", "-T0", f"-{config.compression_level}"] + if config.compression_level > 19: + command.append("--ultra") + proc = await asyncio.create_subprocess_exec( + *command, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE + ) + result, _ = await proc.communicate(uncompressed_data) + if proc.returncode: + raise CalledProcessError(returncode=proc.returncode, cmd=command) + + compressed_data = result + original_size = await zstd_view.resource.get_data_length() + resource.queue_patch(Range(0, original_size), compressed_data) MagicMimeIdentifier.register(ZstdData, "application/x-zstd") diff --git a/ofrak_core/ofrak/ofrak_context.py b/ofrak_core/ofrak/ofrak_context.py index 0b74878e4..8c0ee44d1 100644 --- a/ofrak_core/ofrak/ofrak_context.py +++ b/ofrak_core/ofrak/ofrak_context.py @@ -1,7 +1,7 @@ import asyncio import logging import os -import tempfile +import tempfile312 as tempfile import time from types import ModuleType from typing import Type, Any, Awaitable, Callable, List, Iterable, Optional diff --git a/ofrak_core/ofrak/resource.py b/ofrak_core/ofrak/resource.py index 1d7d3616c..69c651959 100644 --- a/ofrak_core/ofrak/resource.py +++ b/ofrak_core/ofrak/resource.py @@ -4,6 +4,7 @@ import logging from inspect import isawaitable from typing import ( + AsyncIterator, BinaryIO, Iterable, List, @@ -20,6 +21,8 @@ Pattern, overload, ) +from contextlib import asynccontextmanager +import tempfile312 as tempfile from ofrak.component.interface import ComponentInterface from ofrak.model.component_model import ComponentContext, CC, ComponentRunResult @@ -1527,6 +1530,21 @@ async def summarize_tree( tree_string += f"{indent}{branch_symbol}{SPACER_LINE}{child_tree_string}" return tree_string + @asynccontextmanager + async def temp_to_disk( + self, + prefix: Optional[str] = None, + suffix: Optional[str] = None, + dir: Optional[str] = None, + delete: bool = True, + ) -> AsyncIterator[str]: + with tempfile.NamedTemporaryFile( + mode="wb", prefix=prefix, suffix=suffix, dir=dir, delete_on_close=False, delete=delete + ) as temp: + temp.write(await self.get_data()) + temp.close() + yield temp.name + async def save_resources( resources: Iterable["Resource"], diff --git a/ofrak_core/requirements.txt b/ofrak_core/requirements.txt index a65f6c0c5..61c448986 100644 --- a/ofrak_core/requirements.txt +++ b/ofrak_core/requirements.txt @@ -19,6 +19,7 @@ python-magic-bin;platform_system=="Windows" reedsolo==1.7.0 sortedcontainers==2.2.2 synthol~=0.1.1 +tempfile312~=1.0.1 typeguard~=2.13.3 ubi-reader==0.8.5 xattr==0.10.1;platform_system!="Windows" diff --git a/ofrak_core/test_ofrak/components/test_cpio_component.py b/ofrak_core/test_ofrak/components/test_cpio_component.py index 6ea587576..b3a2d9ad5 100644 --- a/ofrak_core/test_ofrak/components/test_cpio_component.py +++ b/ofrak_core/test_ofrak/components/test_cpio_component.py @@ -1,6 +1,7 @@ +import asyncio import os import subprocess -import tempfile +import tempfile312 as tempfile import pytest @@ -26,9 +27,20 @@ async def create_root_resource(self, ofrak_context: OFRAKContext) -> Resource: # Create a CPIO file from the current directory with open(CPIO_ENTRY_NAME, "wb") as f: f.write(INITIAL_DATA) - command = f"find {CPIO_ENTRY_NAME} -print | cpio -o > {TARGET_CPIO_FILE}" - subprocess.run(command, check=True, capture_output=True, shell=True) - result = await ofrak_context.create_root_resource_from_file(TARGET_CPIO_FILE) + cmd = ["cpio", "-o"] + proc = await asyncio.create_subprocess_exec( + *cmd, + cwd=tmpdir, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await proc.communicate(CPIO_ENTRY_NAME.encode()) + if proc.returncode: + raise subprocess.CalledProcessError( + returncode=proc.returncode, cmd=cmd, stdout=stdout, stderr=stderr + ) + result = await ofrak_context.create_root_resource(name=TARGET_CPIO_FILE, data=stdout) os.chdir(wd) return result @@ -46,13 +58,20 @@ async def repack(self, cpio_resource: Resource) -> None: await cpio_resource.pack_recursively() async def verify(self, repacked_cpio_resource: Resource) -> None: - resource_data = await repacked_cpio_resource.get_data() - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(resource_data) - temp_file.flush() - with tempfile.TemporaryDirectory() as temp_flush_dir: - command = f"(cd {temp_flush_dir} && cpio -id < {temp_file.name})" - subprocess.run(command, check=True, capture_output=True, shell=True) - with open(os.path.join(temp_flush_dir, CPIO_ENTRY_NAME), "rb") as f: - patched_data = f.read() - assert patched_data == EXPECTED_DATA + with tempfile.TemporaryDirectory() as temp_flush_dir: + cmd = ["cpio", "-id"] + proc = await asyncio.create_subprocess_exec( + *cmd, + cwd=temp_flush_dir, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await proc.communicate(await repacked_cpio_resource.get_data()) + if proc.returncode: + raise subprocess.CalledProcessError( + returncode=proc.returncode, cmd=cmd, stdout=stdout, stderr=stderr + ) + with open(os.path.join(temp_flush_dir, CPIO_ENTRY_NAME), "rb") as f: + patched_data = f.read() + assert patched_data == EXPECTED_DATA diff --git a/ofrak_core/test_ofrak/components/test_elf_load_aligment_modifier.py b/ofrak_core/test_ofrak/components/test_elf_load_aligment_modifier.py index 35b19ffe7..c73f9b0fa 100644 --- a/ofrak_core/test_ofrak/components/test_elf_load_aligment_modifier.py +++ b/ofrak_core/test_ofrak/components/test_elf_load_aligment_modifier.py @@ -1,6 +1,6 @@ import os import subprocess -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass, field from typing import List, Optional diff --git a/ofrak_core/test_ofrak/components/test_filesystem_component.py b/ofrak_core/test_ofrak/components/test_filesystem_component.py index 88d9fe2a9..12773a2d9 100644 --- a/ofrak_core/test_ofrak/components/test_filesystem_component.py +++ b/ofrak_core/test_ofrak/components/test_filesystem_component.py @@ -4,7 +4,7 @@ import subprocess import sys import warnings -import tempfile +import tempfile312 as tempfile import pytest @@ -328,7 +328,8 @@ def create_broken_symlink(self, root: str): async def create_root_resource(self, ofrak_context: OFRAKContext, directory: str) -> Resource: # Pack with command line `tar` because it supports symbolic links - with tempfile.NamedTemporaryFile(suffix=".tar") as archive: + with tempfile.NamedTemporaryFile(suffix=".tar", delete_on_close=False) as archive: + archive.close() command = ["tar", "--xattrs", "-C", directory, "-cf", archive.name, "."] subprocess.run(command, check=True, capture_output=True) @@ -341,12 +342,8 @@ async def repack(self, root_resource: Resource): await root_resource.pack_recursively() async def extract(self, root_resource: Resource, extract_dir: str): - with tempfile.NamedTemporaryFile(suffix=".tar") as tar: - data = await root_resource.get_data() - tar.write(data) - tar.flush() - - command = ["tar", "--xattrs", "-C", extract_dir, "-xf", tar.name] + async with root_resource.temp_to_disk(suffix=".tar") as tar_path: + command = ["tar", "--xattrs", "-C", extract_dir, "-xf", tar_path] subprocess.run(command, check=True, capture_output=True) diff --git a/ofrak_core/test_ofrak/components/test_general_filesystems.py b/ofrak_core/test_ofrak/components/test_general_filesystems.py index b12b96d11..f39da5084 100644 --- a/ofrak_core/test_ofrak/components/test_general_filesystems.py +++ b/ofrak_core/test_ofrak/components/test_general_filesystems.py @@ -1,5 +1,5 @@ import os -import tempfile +import tempfile312 as tempfile from abc import ABC from typing import Type, List diff --git a/ofrak_core/test_ofrak/components/test_jffs2_component.py b/ofrak_core/test_ofrak/components/test_jffs2_component.py index 818c48349..f97f644e9 100644 --- a/ofrak_core/test_ofrak/components/test_jffs2_component.py +++ b/ofrak_core/test_ofrak/components/test_jffs2_component.py @@ -1,6 +1,6 @@ import os import subprocess -import tempfile +import tempfile312 as tempfile import pytest @@ -44,12 +44,9 @@ async def repack(self, jffs2_resource: Resource) -> None: await jffs2_resource.pack_recursively() async def verify(self, repacked_jffs2_resource: Resource) -> None: - resource_data = await repacked_jffs2_resource.get_data() - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(resource_data) - temp_file.flush() + async with repacked_jffs2_resource.temp_to_disk() as temp_path: with tempfile.TemporaryDirectory() as temp_flush_dir: - command = ["jefferson", "-f", "-d", temp_flush_dir, temp_file.name] + command = ["jefferson", "-f", "-d", temp_flush_dir, temp_path] subprocess.run(command, check=True, capture_output=True) with open(os.path.join(temp_flush_dir, JFFS2_ENTRY_NAME), "rb") as f: patched_data = f.read() diff --git a/ofrak_core/test_ofrak/components/test_lzo_component.py b/ofrak_core/test_ofrak/components/test_lzo_component.py index 4f6d8d983..1249f608c 100644 --- a/ofrak_core/test_ofrak/components/test_lzo_component.py +++ b/ofrak_core/test_ofrak/components/test_lzo_component.py @@ -1,5 +1,4 @@ import subprocess -import tempfile from ofrak.core.lzo import LzoPacker, LzoUnpacker import pytest @@ -26,12 +25,8 @@ def create_test_file(self, tmpdir): self._test_file = compressed_filename async def verify(self, repacked_root_resource: Resource) -> None: - compressed_data = await repacked_root_resource.get_data() - with tempfile.NamedTemporaryFile(suffix=".lzo") as compressed_file: - compressed_file.write(compressed_data) - compressed_file.flush() - - command = ["lzop", "-d", "-f", "-c", compressed_file.name] + async with repacked_root_resource.temp_to_disk() as temp_path: + command = ["lzop", "-d", "-f", "-c", temp_path] result = subprocess.run(command, check=True, capture_output=True) assert result.stdout == self.EXPECTED_REPACKED_DATA diff --git a/ofrak_core/test_ofrak/components/test_patch_symbol_resolution.py b/ofrak_core/test_ofrak/components/test_patch_symbol_resolution.py index 98505e2f7..6e42239cf 100644 --- a/ofrak_core/test_ofrak/components/test_patch_symbol_resolution.py +++ b/ofrak_core/test_ofrak/components/test_patch_symbol_resolution.py @@ -1,5 +1,5 @@ import pytest -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from immutabledict import immutabledict diff --git a/ofrak_core/test_ofrak/components/test_seven_zip_component.py b/ofrak_core/test_ofrak/components/test_seven_zip_component.py index 3c0ee56b0..4e02f9685 100644 --- a/ofrak_core/test_ofrak/components/test_seven_zip_component.py +++ b/ofrak_core/test_ofrak/components/test_seven_zip_component.py @@ -1,6 +1,6 @@ import os import subprocess -import tempfile +import tempfile312 as tempfile import pytest @@ -46,13 +46,9 @@ async def repack(self, seven_zip_resource: Resource) -> None: await seven_zip_resource.pack_recursively() async def verify(self, repacked_seven_zip_resource: Resource) -> None: - resource_data = await repacked_seven_zip_resource.get_data() - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(resource_data) - temp_file.flush() - + async with repacked_seven_zip_resource.temp_to_disk(suffix=".7z") as temp_path: with tempfile.TemporaryDirectory() as temp_flush_dir: - command = ["7zz", "x", f"-o{temp_flush_dir}", temp_file.name] + command = ["7zz", "x", f"-o{temp_flush_dir}", temp_path] subprocess.run(command, check=True, capture_output=True) with open(os.path.join(temp_flush_dir, SEVEN_ZIP_ENTRY_NAME), "rb") as f: patched_data = f.read() diff --git a/ofrak_core/test_ofrak/components/test_squashfs_component.py b/ofrak_core/test_ofrak/components/test_squashfs_component.py index 589182568..6a0fd85fe 100644 --- a/ofrak_core/test_ofrak/components/test_squashfs_component.py +++ b/ofrak_core/test_ofrak/components/test_squashfs_component.py @@ -1,6 +1,6 @@ import os import subprocess -import tempfile +import tempfile312 as tempfile import pytest @@ -43,12 +43,9 @@ async def repack(self, squashfs_resource: Resource) -> None: await squashfs_resource.pack_recursively() async def verify(self, repacked_squashfs_resource: Resource) -> None: - resource_data = await repacked_squashfs_resource.get_data() - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(resource_data) - temp_file.flush() + async with repacked_squashfs_resource.temp_to_disk() as temp_path: with tempfile.TemporaryDirectory() as temp_flush_dir: - command = ["unsquashfs", "-f", "-d", temp_flush_dir, temp_file.name] + command = ["unsquashfs", "-f", "-d", temp_flush_dir, temp_path] subprocess.run(command, check=True, capture_output=True) with open(os.path.join(temp_flush_dir, SQUASH_ENTRY_NAME), "rb") as f: patched_data = f.read() diff --git a/ofrak_core/test_ofrak/components/test_tar_component.py b/ofrak_core/test_ofrak/components/test_tar_component.py index 3e953283c..e42a18b48 100644 --- a/ofrak_core/test_ofrak/components/test_tar_component.py +++ b/ofrak_core/test_ofrak/components/test_tar_component.py @@ -1,6 +1,6 @@ import os import subprocess -import tempfile +import tempfile312 as tempfile import pytest @@ -115,7 +115,8 @@ def setup(self): self.check_stat = False async def create_root_resource(self, ofrak_context: OFRAKContext, directory: str) -> Resource: - with tempfile.NamedTemporaryFile(suffix=".tar") as archive: + with tempfile.NamedTemporaryFile(suffix=".tar", delete_on_close=False) as archive: + archive.close() command = ["tar", "--xattrs", "-C", directory, "-cf", archive.name, "."] subprocess.run(command, check=True, capture_output=True) @@ -128,12 +129,8 @@ async def repack(self, root_resource: Resource): await root_resource.pack_recursively() async def extract(self, root_resource: Resource, extract_dir: str): - with tempfile.NamedTemporaryFile(suffix=".tar") as tar: - data = await root_resource.get_data() - tar.write(data) - tar.flush() - - command = ["tar", "--xattrs", "-C", extract_dir, "-xf", tar.name] + async with root_resource.temp_to_disk(suffix=".tar") as tar_path: + command = ["tar", "--xattrs", "-C", extract_dir, "-xf", tar_path] subprocess.run(command, check=True, capture_output=True) @@ -223,10 +220,6 @@ async def repack(self, root_resource: Resource): await root_resource.pack_recursively() async def extract(self, root_resource: Resource, extract_dir: str): - with tempfile.NamedTemporaryFile(suffix=".tar") as tar: - data = await root_resource.get_data() - tar.write(data) - tar.flush() - - command = ["tar", "--xattrs", "-C", extract_dir, "-xf", tar.name] + async with root_resource.temp_to_disk(suffix=".tar") as tar_path: + command = ["tar", "--xattrs", "-C", extract_dir, "-xf", tar_path] subprocess.run(command, check=True, capture_output=True) diff --git a/ofrak_core/test_ofrak/components/test_ubifs.py b/ofrak_core/test_ofrak/components/test_ubifs.py index 5f86c5899..7d044428e 100644 --- a/ofrak_core/test_ofrak/components/test_ubifs.py +++ b/ofrak_core/test_ofrak/components/test_ubifs.py @@ -1,5 +1,5 @@ import subprocess -import tempfile +import tempfile312 as tempfile import pytest @@ -22,7 +22,8 @@ async def create_root_resource(self, ofrak_context: OFRAKContext, directory: str """ Generated the test UBIFS image with the assistance of the FilesystemPackUnpackVerify test pattern. """ - with tempfile.NamedTemporaryFile() as ubifs_blob: + with tempfile.NamedTemporaryFile(delete_on_close=False) as ubifs_blob: + ubifs_blob.close() command = [ "mkfs.ubifs", "-m", @@ -50,11 +51,7 @@ async def extract(self, root_resource: Resource, extract_dir: str) -> None: expected by the FilesystemPackUnpackVerify pattern. """ - with tempfile.NamedTemporaryFile() as ubifs_blob: - data = await root_resource.get_data() - ubifs_blob.write(data) - ubifs_blob.flush() - - command = ["ubireader_extract_files", "-k", "-o", extract_dir, ubifs_blob.name] + async with root_resource.temp_to_disk() as ubifs_blob_path: + command = ["ubireader_extract_files", "-k", "-o", extract_dir, ubifs_blob_path] subprocess.run(command, check=True, capture_output=True) diff --git a/ofrak_core/test_ofrak/components/test_zstd_component.py b/ofrak_core/test_ofrak/components/test_zstd_component.py index 30e6eae5b..c0b5ea884 100644 --- a/ofrak_core/test_ofrak/components/test_zstd_component.py +++ b/ofrak_core/test_ofrak/components/test_zstd_component.py @@ -1,5 +1,5 @@ import subprocess -import tempfile +import tempfile312 as tempfile from ofrak.core.zstd import ZstdUnpacker, ZstdPacker import pytest @@ -26,13 +26,10 @@ def create_test_file(self, tmpdir): self._test_file = compressed_filename async def verify(self, repacked_root_resource: Resource) -> None: - compressed_data = await repacked_root_resource.get_data() - with tempfile.NamedTemporaryFile(suffix=".zstd") as compressed_file: - compressed_file.write(compressed_data) - compressed_file.flush() + async with repacked_root_resource.temp_to_disk(suffix=".zstd") as compressed_file_path: output_filename = tempfile.mktemp() - command = ["zstd", "-d", "-k", compressed_file.name, "-o", output_filename] + command = ["zstd", "-d", "-k", compressed_file_path, "-o", output_filename] subprocess.run(command, check=True, capture_output=True) with open(output_filename, "rb") as f: result = f.read() diff --git a/ofrak_core/test_ofrak/unit/test_resource.py b/ofrak_core/test_ofrak/unit/test_resource.py index efcb9ea03..445de6a06 100644 --- a/ofrak_core/test_ofrak/unit/test_resource.py +++ b/ofrak_core/test_ofrak/unit/test_resource.py @@ -1,4 +1,4 @@ -import tempfile +import tempfile312 as tempfile from dataclasses import dataclass from io import BytesIO from typing import List, Tuple @@ -209,7 +209,9 @@ async def test_flush_to_disk_pack(ofrak_context: OFRAKContext): # this should not fail because pack_recursively was suppressed await root_resource.write_to(buffer, pack=False) - with tempfile.NamedTemporaryFile() as t: + with tempfile.NamedTemporaryFile(delete_on_close=False) as t: + t.close() + # again, should fail because the packer is run automatically with pytest.raises(MockFailException): await root_resource.flush_data_to_disk(t.name)