Moved 'close' method before asserts when possible

Added unit tests for two fixer methods and unit tests in GitHub

requirements.txt

@@ -1,3 +1,3 @@
-lief==0.12.1
+lief==0.12.3
 python-magic==0.4.27
 yara-python==4.1.0

tests/test_all.py

@@ -2,7 +2,7 @@ import os
 import tempfile
 import unittest
 
-from upxrecoverytool import UpxRecoveryTool, UnsupportedFileError
+from upxrecoverytool import UpxRecoveryTool, UnsupportedFileError, p_info_s
 
 
 class TestInitialChecks(unittest.TestCase):
@@ -41,10 +41,10 @@ class TestInitialChecks(unittest.TestCase):
             urt.init_tmp_buffers()
 
             overlay_size = urt.get_overlay_size()
-            self.assertEqual(overlay_size, 8, f"Wrong detected overlay size {overlay_size} (8 was expected)")
-
             urt.close()
 
+            self.assertEqual(overlay_size, 8, f"Wrong detected overlay size {overlay_size} (8 was expected)")
+
 
 class TestFixes(unittest.TestCase):
 
@@ -52,22 +52,39 @@ class TestFixes(unittest.TestCase):
         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")
 
-            urt.close()
-
-    def test_fix_p_info_filesize(self):
-        pass
-
-    def test_fix_p_info_blocksize(self):
-        pass
-
     def test_fix_p_info_filesize_and_blocksize(self):
-        pass
+
+        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:
@@ -77,6 +94,5 @@ class TestFixes(unittest.TestCase):
             pre_size = os.path.getsize("tests/samples/overlay_8")
             post_size = os.path.getsize(fd.name)
             size_diff = pre_size - post_size
-            self.assertEqual(size_diff, 8, f"Overlay fix error. {size_diff} bytes were removed instead of 8")
-
             urt.close()
+            self.assertEqual(size_diff, 8, f"Overlay fix error. {size_diff} bytes were removed instead of 8")