Modernize the codebase, reformat and improve performance

docs/conf.py

@@ -12,53 +12,54 @@
 #
 import os
 import sys
-sys.path.insert(0, os.path.abspath('..'))
+
+sys.path.insert(0, os.path.abspath(".."))
 
 
 # -- Project information -----------------------------------------------------
 
-project = 'malduck'
-copyright = '2022, CERT Polska'
-author = 'CERT Polska'
+project = "malduck"
+copyright = "2022, CERT Polska"
+author = "CERT Polska"
 
 # The full version, including alpha/beta/rc tags
-version = '4.4.0'
+version = "4.4.0"
 
 # -- General configuration ---------------------------------------------------
 
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    'sphinx.ext.autodoc',
-    'sphinx.ext.viewcode',
-    'sphinx_rtd_theme',
+    "sphinx.ext.autodoc",
+    "sphinx.ext.viewcode",
+    "sphinx_rtd_theme",
 ]
 
 html_theme_options = {
-    'display_version': True,
+    "display_version": True,
 }
 
-project = 'Malduck 🦆'
+project = "Malduck 🦆"
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
-master_doc = 'index'
+master_doc = "index"
 
 # -- Options for HTML output -------------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'sphinx_rtd_theme'
+html_theme = "sphinx_rtd_theme"
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = ["_static"]

malduck/bits.py

@@ -1,6 +1,7 @@
 # Copyright (C) 2018 Jurriaan Bremer.
 # This file is part of Roach - https://github.com/jbremer/roach.
 # See the file 'docs/LICENSE.txt' for copying permission.
+from __future__ import annotations
 
 __all__ = ["rol", "ror", "align", "align_down"]
 

malduck/compression/aplib.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 import struct
 from binascii import crc32
-from typing import Optional
 
 from .components.aplib import APLib
 
@@ -19,18 +20,25 @@ class aPLib:
         from malduck import aplib
 
         # Headerless compressed buffer
-        aplib(b'T\x00he quick\xecb\x0erown\xcef\xaex\x80jumps\xed\xe4veur`t?lazy\xead\xfeg\xc0\x00')
+        aplib(
+            b'T\x00he quick\xecb\x0erown\xcef\xaex\x80'
+            b'jumps\xed\xe4veur`t?lazy\xead\xfeg\xc0\x00'
+        )
         # Header included
-        aplib(b'AP32\x18\x00\x00\x00\r\x00\x00\x00\xbc\x9ab\x9b\x0b\x00\x00\x00\x85\x11J\rh8el\x8eo wnr\xecd\x00')
+        aplib(
+            b'AP32\x18\x00\x00\x00\r\x00\x00\x00\xbc\x9ab'
+            b'\x9b\x0b\x00\x00\x00\x85\x11J\rh8el\x8eo wnr\xecd\x00'
+        )
 
     :param buf: Buffer to decompress
     :type buf: bytes
-    :param headerless: Force headerless decompression (don't perform 'AP32' magic detection)
+    :param headerless:
+        Force headerless decompression (don't perform 'AP32' magic detection)
     :type headerless: bool (default: `True`)
     :rtype: bytes
     """
 
-    def decompress(self, buf: bytes, headerless: bool = True) -> Optional[bytes]:
+    def decompress(self, buf: bytes, headerless: bool = True) -> bytes | None:
         packed_size = None
         packed_crc = None
         orig_size = None

malduck/compression/components/aplib.py

@@ -5,15 +5,16 @@
 Approximately 20 times faster than other Python implementations.
 Compatible with both Python 2 and 3.
 """
+from __future__ import annotations
+
 from io import BytesIO
 
 __all__ = ["APLib"]
 __version__ = "0.6"
 __author__ = "Sandor Nemes"
 
 
-class APLib(object):
-
+class APLib:
     __slots__ = "source", "destination", "tag", "bitcount", "strict"
 
     def __init__(self, source: bytes, strict: bool = True) -> None:
@@ -54,7 +55,6 @@ def depack(self) -> bytes:
         done = False
 
         try:
-
             # first byte verbatim
             self.destination += self.source.read(1)
 
@@ -134,7 +134,10 @@ def pack(self):
 
 def main():
     # self-test
-    data = b"T\x00he quick\xecb\x0erown\xcef\xaex\x80jumps\xed\xe4veur`t?lazy\xead\xfeg\xc0\x00"
+    data = (
+        b"T\x00he quick\xecb\x0erown\xcef\xaex\x80"
+        b"jumps\xed\xe4veur`t?lazy\xead\xfeg\xc0\x00"
+    )
     assert APLib(data).depack() == b"The quick brown fox jumps over the lazy dog"
 
 

malduck/compression/components/lznt1.py

@@ -27,6 +27,8 @@
 https://github.com/libyal/reviveit/
 https://github.com/sleuthkit/sleuthkit/blob/develop/tsk/fs/ntfs.c
 """
+from __future__ import annotations
+
 import array
 import struct
 from io import BytesIO

malduck/compression/gzip.py

@@ -1,14 +1,14 @@
 # Copyright (C) 2018 Jurriaan Bremer.
 # This file is part of Roach - https://github.com/jbremer/roach.
 # See the file 'docs/LICENSE.txt' for copying permission.
-from __future__ import absolute_import
-
-__all__ = ["gzip", "Gzip"]
+from __future__ import annotations
 
 import io
 import zlib
 from gzip import GzipFile
 
+__all__ = ["gzip", "Gzip"]
+
 
 class Gzip:
     r"""

malduck/compression/lznt1.py

@@ -1,11 +1,14 @@
+from __future__ import annotations
+
 from .components.lznt1 import decompress_data
 
 __all__ = ["Lznt1", "lznt1"]
 
 
 class Lznt1:
     """
-    Implementation of LZNT1 decompression. Allows to decompress data compressed by RtlCompressBuffer
+    Implementation of LZNT1 decompression.
+    Allows to decompress data compressed by RtlCompressBuffer
 
     .. code-block:: python
 

malduck/crypto/aes.py

@@ -1,9 +1,9 @@
 # Copyright (C) 2018 Jurriaan Bremer.
 # This file is part of Roach - https://github.com/jbremer/roach.
 # See the file 'docs/LICENSE.txt' for copying permission.
+from __future__ import annotations
 
 import io
-from typing import Optional, Tuple
 
 from Cryptodome.Cipher import AES as AESCipher
 
@@ -27,7 +27,7 @@ class PlaintextKeyBlob(BaseBlob):
 
     def __init__(self) -> None:
         BaseBlob.__init__(self)
-        self.key: Optional[bytes] = None
+        self.key: bytes | None = None
 
     def parse(self, buf: io.BytesIO) -> None:
         """
@@ -42,11 +42,12 @@ def parse(self, buf: io.BytesIO) -> None:
             return
         self.key = value
 
-    def export_key(self) -> Optional[Tuple[str, bytes]]:
+    def export_key(self) -> tuple[str, bytes] | None:
         """
         Exports key from structure or returns None if no key was imported
 
-        :return: Tuple (`algorithm`, `key`). `Algorithm` is one of: "AES-128", "AES-192", "AES-256"
+        :return: Tuple (`algorithm`, `key`).
+            `Algorithm` is one of: "AES-128", "AES-192", "AES-256"
         :rtype: Tuple[str, bytes]
         """
         if self.key is not None:
@@ -164,13 +165,14 @@ class Aes:
     ctr = AesCtr()
 
     @staticmethod
-    def import_key(data: bytes) -> Optional[Tuple[str, bytes]]:
+    def import_key(data: bytes) -> tuple[str, bytes] | None:
         """
         Extracts key from buffer containing :class:`PlaintextKeyBlob` data
 
         :param data: Buffer with `BLOB` structure data
         :type data: bytes
-        :return: Tuple (`algorithm`, `key`). `Algorithm` is one of: "AES-128", "AES-192", "AES-256"
+        :return: Tuple (`algorithm`, `key`).
+            `Algorithm` is one of: "AES-128", "AES-192", "AES-256"
         """
         if len(data) < BLOBHEADER.sizeof():
             return None

malduck/crypto/camellia.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
 __all__ = ["camellia"]

malduck/crypto/chacha20.py

@@ -1,12 +1,12 @@
-from typing import Optional
+from __future__ import annotations
 
 from Cryptodome.Cipher import ChaCha20 as ChaCha20Cipher
 
 __all__ = ["chacha20"]
 
 
 class ChaCha20:
-    def encrypt(self, key: bytes, data: bytes, nonce: Optional[bytes] = None) -> bytes:
+    def encrypt(self, key: bytes, data: bytes, nonce: bytes | None = None) -> bytes:
         """
         Encrypts buffer using ChaCha20 algorithm.
 
@@ -23,7 +23,7 @@ def encrypt(self, key: bytes, data: bytes, nonce: Optional[bytes] = None) -> byt
             nonce = b"\x00" * 8
         return ChaCha20Cipher.new(key=key, nonce=nonce).encrypt(data)
 
-    def decrypt(self, key: bytes, data: bytes, nonce: Optional[bytes] = None) -> bytes:
+    def decrypt(self, key: bytes, data: bytes, nonce: bytes | None = None) -> bytes:
         """
         Decrypts buffer using ChaCha20 algorithm.
 

malduck/crypto/components/pyserpent.py

@@ -34,16 +34,17 @@
 # Anyone thinking of using this code should reconsider. It's slow.
 # Try python-mcrypt instead. In case a faster library is not installed
 # on the target system, this code can be used as a portable fallback.
+from __future__ import annotations
+
 import struct
 import sys
-from typing import List, Optional
 
 block_size = 16
 key_size = 32
 
 
 class Serpent:
-    def __init__(self, key: Optional[bytes] = None) -> None:
+    def __init__(self, key: bytes | None = None) -> None:
         """Serpent."""
 
         if key:
@@ -133,7 +134,7 @@ def byteswap32(x: int) -> int:
     )
 
 
-def set_key(l_key: List[int], key: List[int], key_len: int) -> None:
+def set_key(l_key: list[int], key: list[int], key_len: int) -> None:
     key_len *= 8
     if key_len > 256:
         return None
@@ -962,7 +963,7 @@ def set_key(l_key: List[int], key: List[int], key_len: int) -> None:
     key[4 * 32 + 11] = h
 
 
-def encrypt(key: List[int], in_blk: List[int]) -> None:
+def encrypt(key: list[int], in_blk: list[int]) -> None:
     # serpent_generate.py
     a = in_blk[0]
     b = in_blk[1]
@@ -1946,7 +1947,7 @@ def encrypt(key: List[int], in_blk: List[int]) -> None:
     in_blk[3] = d
 
 
-def decrypt(key: List[int], in_blk: List[int]) -> None:
+def decrypt(key: list[int], in_blk: list[int]) -> None:
     # serpent_generate.py
     a = in_blk[0]
     b = in_blk[1]
@@ -2957,10 +2958,10 @@ def decrypt(key: List[int], in_blk: List[int]) -> None:
 )
 __testdat = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f"
 assert b"\xde&\x9f\xf83\xe42\xb8[.\x88\xd2p\x1c\xe7\\" == Serpent(__testkey).encrypt(
-    __testdat
+    __testdat,
 )
 assert __testdat == Serpent(__testkey).decrypt(
-    b"\xde&\x9f\xf83\xe42\xb8[.\x88\xd2p\x1c\xe7\\"
+    b"\xde&\x9f\xf83\xe42\xb8[.\x88\xd2p\x1c\xe7\\",
 )
 
 

malduck/crypto/des3.py

@@ -1,6 +1,8 @@
 # Copyright (C) 2018 Jurriaan Bremer.
 # This file is part of Roach - https://github.com/jbremer/roach.
 # See the file 'docs/LICENSE.txt' for copying permission.
+from __future__ import annotations
+
 from typing import cast
 
 from Cryptodome.Cipher import DES

malduck/crypto/rabbit.py

@@ -1,9 +1,9 @@
 # Copyright (C) 2018 Jurriaan Bremer.
 # This file is part of Roach - https://github.com/jbremer/roach.
 # See the file 'docs/LICENSE.txt' for copying permission.
+from __future__ import annotations
 
 import struct
-from typing import Optional
 
 from ..bits import rol
 from .xor import xor
@@ -25,7 +25,7 @@ def __init__(self) -> None:
 
 
 class Rabbit:
-    def __init__(self, key: bytes, iv: Optional[bytes]) -> None:
+    def __init__(self, key: bytes, iv: bytes | None) -> None:
         self.ctx = Context()
         self.set_key(key)
         if iv:

malduck/crypto/rc.py

@@ -1,6 +1,7 @@
 # Copyright (C) 2018 Jurriaan Bremer.
 # This file is part of Roach - https://github.com/jbremer/roach.
 # See the file 'docs/LICENSE.txt' for copying permission.
+from __future__ import annotations
 
 from Cryptodome.Cipher import ARC4