From f18f0a116cefab2f4ce9f8fddf00528f876e874c Mon Sep 17 00:00:00 2001
From: dhondta <alexandre.dhondt@gmail.com>
Date: Thu, 16 May 2024 00:41:28 +0200
Subject: [PATCH] Added support for MSDOS Executable (#3)

---
 src/peid/VERSION.txt     |  2 +-
 src/peid/__init__.py     | 10 +++---
 src/peid/db/__init__.py  |  4 +--
 src/peid/exe/__init__.py | 17 +++++++++
 src/peid/exe/msdos.py    | 75 ++++++++++++++++++++++++++++++++++++++++
 src/peid/{ => exe}/pe.py |  5 ++-
 6 files changed, 103 insertions(+), 10 deletions(-)
 create mode 100644 src/peid/exe/__init__.py
 create mode 100644 src/peid/exe/msdos.py
 rename src/peid/{ => exe}/pe.py (96%)

diff --git a/src/peid/VERSION.txt b/src/peid/VERSION.txt
index fd3d3ec..82b225a 100644
--- a/src/peid/VERSION.txt
+++ b/src/peid/VERSION.txt
@@ -1 +1 @@
-2.1.1
+2.2.0
diff --git a/src/peid/__init__.py b/src/peid/__init__.py
index 6ba1d54..1a06416 100644
--- a/src/peid/__init__.py
+++ b/src/peid/__init__.py
@@ -2,7 +2,7 @@
 import os
 
 from .db import SignaturesTree, SignaturesDB
-from .pe import PE
+from .exe import open_exe
 
 __all__ = ["find_ep_only_signature", "identify_packer", "SignaturesDB"]
 
@@ -20,13 +20,15 @@ def find_ep_only_signature(*files, minlength=16, maxlength=64, common_bytes_thre
     data = []
     for f in files:
         try:
-            with PE(f) as pe:
-                bytes_from_ep = [f"{b:02X}" for b in list(pe.read(maxlength, pe.entrypoint_offset))[0]]
+            with open_exe(f) as exe:
+                bytes_from_ep = [f"{b:02X}" for b in list(exe.read(maxlength, exe.entrypoint_offset))[0]]
             maxlength = max(min(len(bytes_from_ep), maxlength), minlength)
             data.append(bytes_from_ep)
-        except (TypeError, ValueError) as e:
+        except (OSError, TypeError, ValueError) as e:
             if logger:
                 logger.debug(f"{f}: {e}")
+            else:
+                raise
     # now determine a signature
     length = maxlength
     while length >= minlength:
diff --git a/src/peid/db/__init__.py b/src/peid/db/__init__.py
index 536fb1f..607a3ba 100644
--- a/src/peid/db/__init__.py
+++ b/src/peid/db/__init__.py
@@ -2,7 +2,7 @@
 import re
 from os.path import abspath, basename, dirname, exists, expanduser, join
 
-from ..pe import PE
+from ..exe import open_exe
 
 
 __all__ = ["DB", "SignaturesDB", "SignaturesTree"]
@@ -77,7 +77,7 @@ def _match(subtree, byteseq):
                     subtree = subtree[byte]
                 else:
                     break
-        with PE(pe, logger=self.logger) as f:
+        with open_exe(pe, logger=self.logger) as f:
             if ep_only:
                 for byteseq in f.read(n_bytes, f.entrypoint_offset):
                     _match(self.__tree['ep_only'], byteseq)
diff --git a/src/peid/exe/__init__.py b/src/peid/exe/__init__.py
new file mode 100644
index 0000000..25f9f33
--- /dev/null
+++ b/src/peid/exe/__init__.py
@@ -0,0 +1,17 @@
+# -*- coding: UTF-8 -*-
+from .msdos import MSDOS
+from .pe import PE
+
+
+__all__ = ["open_exe", "MSDOS", "PE"]
+
+
+def open_exe(path, logger=None):
+    """ Find a matching format and return the instantiated executable object. """
+    for fmt in [PE, MSDOS]:
+        try:
+            return fmt(path, logger)
+        except OSError:
+            pass
+    raise OSError("Not a valid executable or supported executable format")
+
diff --git a/src/peid/exe/msdos.py b/src/peid/exe/msdos.py
new file mode 100644
index 0000000..4e1f4e4
--- /dev/null
+++ b/src/peid/exe/msdos.py
@@ -0,0 +1,75 @@
+# -*- coding: UTF-8 -*-
+import struct
+from os.path import getsize
+
+
+__all__ = ["MSDOS"]
+
+
+class MSDOS:
+    def __init__(self, path, logger=None):
+        self.path, self.size, self.logger = path, getsize(path), logger
+        self.__fd = f = open(path, "rb")
+        h = f.read(64)
+        # check MZ signature
+        if h[:2] != b"MZ":
+            raise OSError("Invalid MZ signature")
+        # read some header fields
+        self.bytes_last_page = struct.unpack("<H", h[2:4])[0] or 512
+        self.pages_in_file = struct.unpack("<H", h[4:6])[0]
+        self.number_relocations = struct.unpack("<H", h[6:8])[0]
+        self.header_paragraphs = struct.unpack("<H", h[8:10])[0]
+        self.initial_ip = struct.unpack("<H", h[20:22])[0]
+        self.initial_cs = struct.unpack("<H", h[22:24])[0]
+        self.relocation_table_offset = struct.unpack("<H", h[24:26])[0]
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, type, value, traceback):
+        self.close()
+    
+    def close(self):
+        self.__fd.close()
+    
+    def itersections(self):
+        f = self.__fd
+        for i in range(self.number_of_sections):
+            f.seek(self.pe_offset + 24 + self.size_of_opt_header + i * 40 + 8)
+            virtual_size = int.from_bytes(f.read(4), "little")
+            virtual_addr = int.from_bytes(f.read(4), "little")
+            raw_size     = int.from_bytes(f.read(4), "little")
+            raw_pointer  = int.from_bytes(f.read(4), "little")
+            yield virtual_size, virtual_addr, raw_size, raw_pointer
+    
+    def read(self, n=64, *offsets):
+        if len(offsets) == 0:
+            offsets = range(0, self.size-n)
+        for o in offsets:
+            self.__fd.seek(o)
+            r = self.__fd.read(min(n, self.size-o))
+            if self.logger:
+                self.logger.debug(" ".join(f"{b:02X}" for b in r))
+            yield r
+    
+    @property
+    def entrypoint_offset(self):
+        return (self.initial_cs << 4) + self.initial_ip
+    
+    @property
+    def file_size(self):
+        return (self.pages_in_file - 1) * 512 + self.bytes_last_page
+    
+    @property
+    def header_size(self):
+        return self.header_paragraphs * 16
+    
+    @property
+    def sections_offsets(self):
+        f, offsets = self.__fd, []
+        f.seek(self.relocation_table_offset)
+        for _ in range(self.number_relocations):
+            segment, offset = struct.unpack("<HH", f.read(4))
+            offsets.append((segment << 4) + offset)
+        return offsets
+
diff --git a/src/peid/pe.py b/src/peid/exe/pe.py
similarity index 96%
rename from src/peid/pe.py
rename to src/peid/exe/pe.py
index 0f519e0..df03f5c 100644
--- a/src/peid/pe.py
+++ b/src/peid/exe/pe.py
@@ -1,6 +1,5 @@
 # -*- coding: UTF-8 -*-
 import builtins
-from functools import lru_cache, wraps
 from os.path import getsize
 
 
@@ -26,10 +25,10 @@ def __init__(self, path, logger=None):
         f.seek(self.pe_offset)
         if f.read(4) != b"PE\x00\x00":
             raise OSError("Invalid PE signature")
-        # ready NumberOfSections 
+        # read NumberOfSections 
         f.seek(self.pe_offset + 6)
         self.number_of_sections = int.from_bytes(f.read(2), "little")
-        # ready SizeOfOptionalHeader
+        # read SizeOfOptionalHeader
         f.seek(self.pe_offset + 20)
         self.size_of_opt_header = int.from_bytes(f.read(2), "little")