From d50ec846c2b87ab127fb74baff210aff5f23d6b1 Mon Sep 17 00:00:00 2001
From: nganhkhoa <mail.nganhkhoa@gmail.com>
Date: Fri, 30 Aug 2024 21:23:14 +0700
Subject: [PATCH] [python] some updates

- add ELF
- add basic ROMFS
---
 python/main.py             |  5 ++++-
 python/matcher/__init__.py |  5 +++++
 python/matcher/elf.py      | 17 +++++++++++++++++
 python/matcher/romfs.py    | 39 ++++++++++++++++++++++++++++++++++++++
 python/matcher/zip.py      |  8 +++++---
 5 files changed, 70 insertions(+), 4 deletions(-)
 create mode 100644 python/matcher/elf.py
 create mode 100644 python/matcher/romfs.py

diff --git a/python/main.py b/python/main.py
index f07ac9a..b24a272 100644
--- a/python/main.py
+++ b/python/main.py
@@ -8,7 +8,10 @@ signatures = [
     matcher.Zip,
     matcher.Ambarella,
     matcher.SquashFS,
-    matcher.FlattenDeviceTree
+    matcher.RomFS,
+    matcher.CromFS,
+    matcher.FlattenDeviceTree,
+    matcher.ELF,
 ]
 
 def detect(args):
diff --git a/python/matcher/__init__.py b/python/matcher/__init__.py
index f4c98f7..ba89eea 100644
--- a/python/matcher/__init__.py
+++ b/python/matcher/__init__.py
@@ -11,3 +11,8 @@ from .flatten_device_tree import FlattenDeviceTree
 # file system formats
 from .squashfs import SquashFS
 from .ubifs import UbiFS
+from .romfs import RomFS
+from .cromfs import CromFS
+
+# common executable formats
+from .elf import ELF
diff --git a/python/matcher/elf.py b/python/matcher/elf.py
new file mode 100644
index 0000000..cbd8595
--- /dev/null
+++ b/python/matcher/elf.py
@@ -0,0 +1,17 @@
+import io
+
+from .matcher import SignatureMatcher, Match
+
+class ELF(SignatureMatcher):
+    def __init__(self, file):
+        self.name = "ELF"
+        self.signature = b'\x7f\x45\x4c\x46\x02\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00'
+        super().__init__(file)
+
+    def is_valid(self):
+        for match in self.search():
+            # walk back for the firmware section header
+            start = match
+            self.matches += [Match(start, 0)]
+
+        return len(self.matches) != 0
diff --git a/python/matcher/romfs.py b/python/matcher/romfs.py
new file mode 100644
index 0000000..b6ae413
--- /dev/null
+++ b/python/matcher/romfs.py
@@ -0,0 +1,39 @@
+import io
+
+from PyRomfsImage import *
+
+from .matcher import SignatureMatcher, Match
+
+class RomFS(SignatureMatcher):
+    """
+    RomFS a readonly file system
+    big-endian
+    https://www.kernel.org/doc/Documentation/filesystems/romfs.txt
+    """
+    def __init__(self, file):
+        self.name = "RomFS"
+        self.signature = b'-rom1fs-'
+        super().__init__(file)
+
+    def is_valid(self):
+        for match in self.search():
+            start = match
+            header = io.BytesIO(self.file[start:start+14])
+            magic = header.read(8)
+            fullsize = header.read(4)
+            checksum = header.read(4)
+            fullsize = int.from_bytes(fullsize, 'big')
+            self.matches += [Match(start, fullsize)]
+
+        return len(self.matches) != 0
+
+    def view(self, match):
+
+        start = match.offset
+        length = match.length
+        raw = io.BytesIO(self.file[start:start+length])
+        rom = Romfs(raw)
+        root = rom.getRoot()
+        print(root.name)
+        for ch in root.children:
+            print(ch.name)
diff --git a/python/matcher/zip.py b/python/matcher/zip.py
index b561c96..0a91358 100644
--- a/python/matcher/zip.py
+++ b/python/matcher/zip.py
@@ -29,9 +29,11 @@ class Zip(SignatureMatcher):
             file_name_length = header.read(2)
             extra_field_length = header.read(2)
 
-            file_name_length = int.from_bytes(file_name_length, 'little')
-            extra_field_length = int.from_bytes(extra_field_length, 'little')
-            compressed_size = int.from_bytes(compressed_size, 'little')
+            as_num = lambda x: int.form_bytes(x, 'little')
+
+            file_name_length = as_num(file_name_length)
+            extra_field_length = as_num(extra_field_length)
+            compressed_size = as_num(compressed_size)
 
             header_size = 4*4 + 2*7
             data = {