Merge pull request #11 from NozomiNetworks/p_info_tests

p_info_tests

.github/workflows/pytest.yaml

@@ -0,0 +1,28 @@
+name: Unit tests
+
+on:
+  [push]
+
+jobs:
+  run_tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.x
+        uses: actions/setup-python@v4
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.x'
+
+      # You can test your matrix by printing the current Python version
+      # - name: Display Python version
+      #  run: python -c "import sys; print(sys.version)"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+      - name: Test with pytest
+        run: |
+          pip install pytest
+          pytest tests/test_all.py --doctest-modules --junitxml=junit/test-results.xml

.gitignore

@@ -0,0 +1,2 @@
+.vscode
+__pycache__

README.md

@@ -8,10 +8,12 @@ files done to prevent their automatic unpacking with a standard UPX tool.
 This tool detects and repairs the following common modifications:
 - `l_magic` field of the `l_info` structure (`UPX!` magic value)
 - `p_filesize` and `p_blocksize` fields of the `p_info` structure
+- Overlay bytes
 
 ### Dependencies
 
 The script requires the following libraries listed on `requirements.txt`:
 
-- [`elftools`](https://github.com/eliben/pyelftools)
+- [`lief`](https://lief-project.github.io)
 - [`python-magic`](https://pypi.org/project/python-magic/)
+- [`yara-python`](https://github.com/VirusTotal/yara-python)

requirements.txt

@@ -1,2 +1,3 @@
-pyelftools==0.28
-python-magic==0.4.27
+lief==0.12.3
+python-magic==0.4.27
+yara-python==4.1.0

rules/arm.yar

@@ -0,0 +1,35 @@
+import "elf"
+
+rule upx_entry_point {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file is packed with UPX based on its entry point."
+		
+	strings:
+		$ep_arm_0 = {1c c0 4f e2 06 4c 9c e8 02 00 a0 e1 0c b0 8b e0 0c a0 8a e0 00 30 9b e5 01 90 4c e0 01 20 a0 e1}
+		$ep_arm_1 = {18 d0 4d e2 ?? 02 00 eb 00 c0 dd e5 0e 00 5c e3 ?? 02 00 1a 0c 48 2d e9 00 b0 d0 e5 06 cc a0 e3}
+		$ep_arm_2 = {00 18 d0 4d e2 9c 00 00 eb 00 10 81 e0 3e 40 2d e9 00 50 e0 e3 02 41 a0 e3 19 00 00 ea 1a 00 bd}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and for any of ($ep_*):($ at elf.entry_point)
+}
+
+rule upx_init_code_not_ep {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file contains UPX code that should be at the entry point but its located anywhere else."
+		
+	strings:
+		$ep_arm_0 = {1c c0 4f e2 06 4c 9c e8 02 00 a0 e1 0c b0 8b e0 0c a0 8a e0 00 30 9b e5 01 90 4c e0 01 20 a0 e1}
+		$ep_arm_1 = {18 d0 4d e2 ?? 02 00 eb 00 c0 dd e5 0e 00 5c e3 ?? 02 00 1a 0c 48 2d e9 00 b0 d0 e5 06 cc a0 e3}
+		$ep_arm_2 = {00 18 d0 4d e2 9c 00 00 eb 00 10 81 e0 3e 40 2d e9 00 50 e0 e3 02 41 a0 e3 19 00 00 ea 1a 00 bd}
+
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and any of ($ep_*)
+        and for 0 of ($ep_*):($ at elf.entry_point)
+}

rules/intel_80386.yar

@@ -0,0 +1,31 @@
+import "elf"
+
+rule upx_entry_point {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file is packed with UPX based on its entry point."
+		
+	strings:
+		$ep_x86 = {50 E8 [4] EB 0E 5A 58 59 97 60 8A 54 24 20 E9 [4] 60}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and for any of ($ep_*):($ at elf.entry_point)
+}
+
+rule upx_init_code_not_ep {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file contains UPX code that should be at the entry point but its located anywhere else."
+		
+	strings:
+		$ep_x86 = {50 E8 [4] EB 0E 5A 58 59 97 60 8A 54 24 20 E9 [4] 60}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and any of ($ep_*)
+        and for 0 of ($ep_*):($ at elf.entry_point)
+}
+

rules/mips.yar

@@ -0,0 +1,36 @@
+import "elf"
+
+rule upx_entry_point {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file is packed with UPX based on its entry point."
+		
+	strings:
+		$ep_mips_be_1 = {04 11 [2] 27 fe 00 00 27 bd ff fc af bf 00 00 00 a4 28 20 ac e6 00 00 3c 0d 80 00 01 a0 48 21 24 0b 00 01 04 11}
+		$ep_mips_be_2 = {04 11 [2] 27 f7 00 00 90 99 00 00 24 01 fa 00 90 98 00 01 33 22 00 07 00 19 c8 c2 03 21 08 04}
+		$ep_mips_le_1 = {?? ?? 11 04 00 00 fe 27 fc ff bd 27 00 00 bf af 20 28 a4 00 00 00 e6 ac 00 80 0d 3c 21 48 a0 01 01 00 0b 24 [2] 11 04}
+		$ep_mips_le_2 = {?? ?? 11 04 00 00 f7 27 00 00 99 90 00 fa 01 24 01 00 98 90 07 00 22 33 c2 c8 19 00 04 08 21 03}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and for any of ($ep_*):($ at elf.entry_point)
+}
+
+rule upx_init_code_not_ep {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file contains UPX code that should be at the entry point but its located anywhere else."
+		
+	strings:
+		$ep_mips_be_1 = {04 11 [2] 27 fe 00 00 27 bd ff fc af bf 00 00 00 a4 28 20 ac e6 00 00 3c 0d 80 00 01 a0 48 21 24 0b 00 01 04 11}
+		$ep_mips_be_2 = {04 11 [2] 27 f7 00 00 90 99 00 00 24 01 fa 00 90 98 00 01 33 22 00 07 00 19 c8 c2 03 21 08 04}
+		$ep_mips_le_1 = {?? ?? 11 04 00 00 fe 27 fc ff bd 27 00 00 bf af 20 28 a4 00 00 00 e6 ac 00 80 0d 3c 21 48 a0 01 01 00 0b 24 [2] 11 04}
+		$ep_mips_le_2 = {?? ?? 11 04 00 00 f7 27 00 00 99 90 00 fa 01 24 01 00 98 90 07 00 22 33 c2 c8 19 00 04 08 21 03}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and any of ($ep_*)
+        and for 0 of ($ep_*):($ at elf.entry_point)
+}

rules/powerpc.yar

@@ -0,0 +1,32 @@
+import "elf"
+
+rule upx_entry_point {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file is packed with UPX based on its entry point."
+		
+	strings:
+		$ep_ppc_0 = {48 00 ?? ?? 7c 00 29 ec 7d a8 02 a6 28 07 00 02 40 82 00 e4 90 a6 00 00 }
+		$ep_ppc_1 = {48 00 ?? ?? 28 07 00 0e 40 82 0a 4c 94 21 ff e8 7c 08 02 a6 7c c9 33 78 81 06 00 00 7c a7 2b 78}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and for any of ($ep_*):($ at elf.entry_point)
+}
+
+rule upx_init_code_not_ep {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file contains UPX code that should be at the entry point but its located anywhere else."
+		
+	strings:
+		$ep_ppc_0 = {48 00 ?? ?? 7c 00 29 ec 7d a8 02 a6 28 07 00 02 40 82 00 e4 90 a6 00 00 }
+		$ep_ppc_1 = {48 00 ?? ?? 28 07 00 0e 40 82 0a 4c 94 21 ff e8 7c 08 02 a6 7c c9 33 78 81 06 00 00 7c a7 2b 78}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and any of ($ep_*)
+        and for 0 of ($ep_*):($ at elf.entry_point)
+}

rules/x86-64.yar

@@ -0,0 +1,32 @@
+import "elf"
+
+rule upx_entry_point {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file is packed with UPX based on its entry point."
+		
+	strings:
+		$ep_x64_0 = {50 52 e8 [4] 55 53 51 52 48 01 fe 56 48 89 fe 48 89 d7 31 db 31 c9 48 83 cd ff e8}
+		$ep_x64_1 = {50 52 e8 [4] 55 53 51 52 48 01 fe 56 41 80 f8 0e 0f [5] 55 48 89 e5 44 8b 09}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and for any of ($ep_*):($ at elf.entry_point)
+}
+
+rule upx_init_code_not_ep {
+
+	meta:
+		author = "Nozomi Networks Labs"
+		description = "Rule to detect if an ELF file contains UPX code that should be at the entry point but its located anywhere else."
+		
+	strings:
+		$ep_x64_0 = {50 52 e8 [4] 55 53 51 52 48 01 fe 56 48 89 fe 48 89 d7 31 db 31 c9 48 83 cd ff e8}
+		$ep_x64_1 = {50 52 e8 [4] 55 53 51 52 48 01 fe 56 41 80 f8 0e 0f [5] 55 48 89 e5 44 8b 09}
+
+	condition:
+		uint32(0)==0x464c457f // ELF header
+        and any of ($ep_*)
+        and for 0 of ($ep_*):($ at elf.entry_point)
+}

tests/test_all.py

@@ -0,0 +1,98 @@
+import os
+import tempfile
+import unittest
+
+from upxrecoverytool import UpxRecoveryTool, UnsupportedFileError, p_info_s
+
+
+class TestInitialChecks(unittest.TestCase):
+
+    def test_check_file_type(self):
+
+        with tempfile.NamedTemporaryFile() as fd1:
+            with tempfile.NamedTemporaryFile() as fd2:
+                with self.assertRaises(UnsupportedFileError):
+                    UpxRecoveryTool(fd1.name, fd2.name, False)
+
+    def test_is_upx(self):
+
+        # UPX file
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/overlay_8", fd.name, False)
+            self.assertTrue(urt.is_upx(), "File not detected as UPX compressed")
+            urt.close()
+
+        # Non-UPX file
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/no_upx", fd.name, True)
+            self.assertFalse(urt.is_upx(), "File detected as UPX compressed")
+            urt.close()
+
+        # UPX with hidden real EP
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/hidden_ep", fd.name, False)
+            self.assertTrue(urt.is_upx(), "Hidden UPX EP was not detected")
+            urt.close()
+
+    def test_get_overlay_size(self):
+
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/overlay_8", fd.name, False)
+            urt.init_tmp_buffers()
+
+            overlay_size = urt.get_overlay_size()
+            urt.close()
+
+            self.assertEqual(overlay_size, 8, f"Wrong detected overlay size {overlay_size} (8 was expected)")
+
+
+class TestFixes(unittest.TestCase):
+
+    def test_fix_l_info(self):
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/l_info", fd.name, False)
+            urt.fix()
+            urt.close()
+
+            for upx_sig_off in [0xEC, 0x1C1B, 0x2403, 0x240C]:
+                fd.seek(upx_sig_off, os.SEEK_SET)
+                sig = fd.read(4)
+                self.assertEqual(sig, b"UPX!", f"UPX! sig at 0x{upx_sig_off:X} wasn't fixed")
+
+    def test_fix_p_info_filesize_and_blocksize(self):
+
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/p_info_0", fd.name, False)
+            urt.fix()
+
+            fd.seek(0xf4, os.SEEK_SET)
+            fixed_p_info = p_info_s(fd.read(12))
+            urt.close()
+            self.assertEqual(fixed_p_info.p_blocksize, b"\x38\x49\x00\x00", f"Error fixing p_info.p_blocksize. \
+                                38490000 expected but {fixed_p_info.p_blocksize.hex()} was read")
+            self.assertEqual(fixed_p_info.p_filesize, b"\x38\x49\x00\x00", f"Error fixing p_info.p_blocksize. \
+                                38490000 expected but {fixed_p_info.p_filesize.hex()} was read")
+
+    def test_fix_p_info_sizes_mismatch(self):
+
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/p_info_size_mismatch", fd.name, False)
+            urt.fix()
+
+            fd.seek(0xf4, os.SEEK_SET)
+            fixed_p_info = p_info_s(fd.read(12))
+            urt.close()
+            self.assertEqual(fixed_p_info.p_blocksize, fixed_p_info.p_filesize,
+                             (f"Error fixing p_info structure. p_blocksize ({fixed_p_info.p_blocksize.hex()}) and "
+                              f"p_filesize ({fixed_p_info.p_filesize.hex()})"))
+
+    def test_fix_overlay(self):
+        with tempfile.NamedTemporaryFile() as fd:
+            urt = UpxRecoveryTool("tests/samples/overlay_8", fd.name, False)
+            urt.fix()
+
+            pre_size = os.path.getsize("tests/samples/overlay_8")
+            post_size = os.path.getsize(fd.name)
+            size_diff = pre_size - post_size
+            urt.close()
+            self.assertEqual(size_diff, 8, f"Overlay fix error. {size_diff} bytes were removed instead of 8")

upxrecoverytool.py

@@ -1,14 +1,14 @@
 #! /usr/bin/python3
 
 import os
-import re
+import lief
 import mmap
+import yara
 import magic
 import shutil
 import struct
 import argparse
 import tempfile
-from elftools.elf.elffile import ELFFile
 
 
 class UnsupportedFileError(Exception):
@@ -71,37 +71,14 @@ class PackHeader:
 class UpxRecoveryTool:
 
     upx_sigs = {
-        "Intel 80386": [
-            br"^\x50\xe8.{2}\x00\x00\xeb.\x5a\x58\x59\x97\x60\x8a\x54\x24\x20\xe9.{2}\x00\x00\x60\x8b\x74\x24\x24\x8b\x7c\x24\x2c\x83\xcd",
-        ],
-        "x86-64": [
-            # ucl
-            br"^\x50\x52\xe8.{4}\x55\x53\x51\x52\x48\x01\xfe\x56\x48\x89\xfe\x48\x89\xd7\x31\xdb\x31\xc9\x48\x83\xcd\xff\xe8",
-            # lzma
-            br"^\x50\x52\xe8.{4}\x55\x53\x51\x52\x48\x01\xfe\x56\x41\x80\xf8\x0e\x0f.{5}\x55\x48\x89\xe5\x44\x8b\x09",
-        ],
-        "MIPS": [
-            # upx_mips_be
-            br"^\x04\x11.{2}\x27\xfe\x00\x00\x27\xbd\xff\xfc\xaf\xbf\x00\x00\x00\xa4\x28\x20\xac\xe6\x00\x00\x3c\x0d\x80\x00\x01\xa0\x48\x21\x24\x0b\x00\x01\x04\x11",
-            br"^\x04\x11.{2}\x27\xf7\x00\x00\x90\x99\x00\x00\x24\x01\xfa\x00\x90\x98\x00\x01\x33\x22\x00\x07\x00\x19\xc8\xc2\x03\x21\x08\x04",
-            # upx_mips_le
-            br"^.{2}\x11\x04\x00\x00\xfe\x27\xfc\xff\xbd\x27\x00\x00\xbf\xaf\x20\x28\xa4\x00\x00\x00\xe6\xac\x00\x80\x0d\x3c\x21\x48\xa0\x01\x01\x00\x0b\x24.{2}\x11\x04",
-            br"^.{2}\x11\x04\x00\x00\xf7\x27\x00\x00\x99\x90\x00\xfa\x01\x24\x01\x00\x98\x90\x07\x00\x22\x33\xc2\xc8\x19\x00\x04\x08\x21\x03",
-        ],
-        "ARM": [
-            br"^\x1c\xc0\x4f\xe2.{2}\x9c\xe8\x02\x00\xa0\xe1\x0c\xb0\x8b\xe0\x0c\xa0\x8a\xe0\x00\x30\x9b\xe5\x01\x90\x4c\xe0\x01\x20\xa0\xe1",
-            br"^\x18\xd0\x4d\xe2.{2}\x00\xeb\x00\xc0\xdd\xe5\x0e\x00\x5c\xe3.\x02\x00\x1a\x0c\x48\x2d\xe9\x00\xb0\xd0\xe5\x06\xcc\xa0\xe3",
-            br"^\x18\xd0\x4d\xe2.{2}\x00\xeb\x00\x10\x81\xe0\x3e\x40\x2d\xe9\x00\x50\xe0\xe3\x02\x41\xa0\xe3.{2}\x00\xea\x1a\x00\xbd\xe8",
-        ],
-        "PowerPC or cisco 4500": [
-            # ucl
-            br"^\x48\x00.{2}\x7c\x00\x29\xec\x7d\xa8\x02\xa6\x28\x07\x00\x02\x40\x82\x00\xe4\x90\xa6\x00\x00",
-            # lzma
-            br"^\x48\x00.{2}\x28\x07\x00\x0e\x40\x82\x0a\x4c\x94\x21\xff\xe8\x7c\x08\x02\xa6\x7c\xc9\x33\x78\x81\x06\x00\x00\x7c\xa7\x2b\x78",
-        ],
+        "Intel 80386": "intel_80386.yar",
+        "x86-64": "x86-64.yar",
+        "MIPS": "mips.yar",
+        "ARM": "arm.yar",
+        "PowerPC or cisco 4500": "powerpc.yar",
     }
 
-    def __init__(self, in_file, out_file):
+    def __init__(self, in_file, out_file, assume_upx):
         """ Initialization method. Receives the path to the file to be fixed and the output path for the result """
 
         self.in_fd = None
@@ -117,14 +94,17 @@ class UpxRecoveryTool:
 
         # File is ELF, so it can be parsed
         self.in_fd = open(self.in_file, "rb")
-        self.elf = ELFFile(self.in_fd)
+        self.elf = lief.parse(self.in_file)
 
         # Get file size for boudaries checks
         self.file_size = os.fstat(self.in_fd.fileno()).st_size
 
-        # Check if it is packed with UPX
-        if not self.is_upx():
-            raise NonUpxError
+        if not assume_upx:
+            # Check if it is packed with UPX
+            if not self.is_upx():
+                raise NonUpxError
+        else:
+            print("[i] Assuming file is UPX")
 
         # Get UPX version. p_info fix doesn't work with UPX 4
         self.detect_version()
@@ -151,14 +131,23 @@ class UpxRecoveryTool:
     def is_upx(self):
         """ Method that looks for ASM signatures to identify a UPX executable """
 
-        # Instead of looking for the pattern in all the executable bytes, look for the
-        # UPX sigs at the EP. The code is not so beautiful but it'll be more efficient.
-        ep_bytes = self.get_ep_bytes(50)
+        rules_path = os.path.join("rules", self.upx_sigs[self.arch])
+        rules = yara.compile(rules_path)
+        matches = rules.match(self.in_file)
 
-        for sig in self.upx_sigs[self.arch]:
-            if re.match(sig, ep_bytes, re.DOTALL):
-                print("[i] File is UPX")
-                return True
+        if matches:
+            print("[i] File is UPX")
+
+            if len(matches) == 1:
+                if matches[0].rule == "upx_init_code_not_ep":
+                    print(
+                        "[!] UPX entry point signature is in an address different from the entry point")
+
+            # 2 matches. EP code found in EP and other address
+            else:
+                print("[!] Multiple UPX entrypoint code found in the same file")
+
+            return True
 
         return False
 
@@ -184,7 +173,7 @@ class UpxRecoveryTool:
         eh = self.elf.header
 
         # l_info
-        self.l_info_off = eh.e_phoff + eh.e_phnum * eh.e_phentsize
+        self.l_info_off = eh.program_header_offset + eh.numberof_segments * eh.program_header_size
 
         if self.l_info_off + 12 > self.file_size:
             raise CorruptedFileError("Parsing l_info structure")
@@ -212,26 +201,33 @@ class UpxRecoveryTool:
     def get_ep_bytes(self, num_bytes):
         """ Method to get the first 'num_bytes' of the executable's Entry Point. Used to apply UPX signatures """
 
-        ep_bytes = None
-        ep = self.elf.header['e_entry']
-
-        # Loop segments looking where the EP is
-        for seg in self.elf.iter_segments():
-            sh = seg.header
-            if ep > sh.p_vaddr and ep < sh.p_vaddr + sh.p_memsz:
-                start_off = ep - sh.p_vaddr
-
-                if start_off + num_bytes > self.file_size:
-                    raise CorruptedFileError("Walking through the program headers")
-
-                ep_bytes = seg.data()[start_off: start_off + num_bytes]
+        ep_bytes_list = self.elf.get_content_from_virtual_address(self.elf.entrypoint, num_bytes)
+        ep_bytes = bytearray(ep_bytes_list)
 
         return ep_bytes
 
-    def fix(self):
-        """ Method to fix all the (supported) modifications of UPX """
+    def get_overlay_size(self):
+        """ Method to check if it seems that the file contains overlay data after a proper PackHeader """
 
-        fixed = False
+        overlay_size = 0
+        upx_count = 0
+
+        upx_offset = self.buff.find(b"UPX!", 0)
+
+        while upx_offset != -1:
+            upx_count += 1
+            last_upx_offset = upx_offset
+            upx_offset = self.buff.find(b"UPX!", upx_offset + 4)
+
+        # If there are less than 3 UPX! sigs, we can't be sure the PackHeader can be easily found
+        if upx_count >= 3:
+            if last_upx_offset < self.buff.size() - 36:
+                overlay_size = self.buff.size() - 36 - last_upx_offset
+
+        return overlay_size
+
+    def init_tmp_buffers(self):
+        """ Method to initialize internal temporary buffers """
 
         self.tmp_folder = tempfile.TemporaryDirectory()
         self.tmp_file = os.path.join(self.tmp_folder.name, os.path.basename(self.out_file))
@@ -240,6 +236,12 @@ class UpxRecoveryTool:
         self.tmp_fd = open(self.tmp_file, "r+b")
         self.buff = mmap.mmap(self.tmp_fd.fileno(), 0)
 
+    def fix(self):
+        """ Method to fix all the (supported) modifications of UPX """
+
+        fixed = False
+
+        self.init_tmp_buffers()
         self.load_structs()
 
         fixed |= self.fix_l_info()
@@ -247,6 +249,8 @@ class UpxRecoveryTool:
         # Now that UPX! magic bytes are restored, PackHeader can be properly loaded
         self.pack_hdr = PackHeader(self.buff)
 
+        fixed |= self.fix_overlay()
+
         if self.version != 4:
             fixed |= self.fix_p_info()
 
@@ -328,6 +332,23 @@ class UpxRecoveryTool:
         int_size = struct.unpack("<i", self.pack_hdr.u_file_size)[0]
         print(f"  [i] Fixed p_info sizes with value 0x{int_size:x} from PackHeader")
 
+    def fix_overlay(self):
+        """ Method to crop the file to remove overlay bytes  """
+
+        fixed = False
+
+        overlay_size = self.get_overlay_size()
+        if overlay_size > 0:
+            new_size = self.buff.size() - overlay_size
+            print(f"[i] Removing {overlay_size} bytes of overlay")
+
+            # self.buff.resize may fail on BSD and Mac
+            self.tmp_fd.truncate(new_size)
+
+            fixed = True
+
+        return fixed
+
     def close(self):
         """ Method to close memory buffers and file descriptors """
         # Close mmap
@@ -343,7 +364,8 @@ class UpxRecoveryTool:
             self.tmp_fd.close()
 
         # Remove temporary file and dir
-        self.tmp_folder.cleanup()
+        if self.tmp_folder is not None:
+            self.tmp_folder.cleanup()
 
 
 if __name__ == "__main__":
@@ -351,12 +373,14 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Script to check and fix UPX files modifications")
     parser.add_argument('-i', dest='input', required=True, help="Path to supposed UPX file to be fixed")
     parser.add_argument('-o', dest='output', required=True, help="Path to write the fixed version of the file")
+    parser.add_argument('-a', '--assume-upx', action='store_true', help=f"Assume file is UPX. Use it when \
+        {parser.prog} doesn't detect the input as UPX and you think it is wrong.")
     args = parser.parse_args()
 
     urt = None
 
     try:
-        urt = UpxRecoveryTool(args.input, args.output)
+        urt = UpxRecoveryTool(args.input, args.output, args.assume_upx)
         urt.fix()
 
     except UnsupportedFileError as why: