From e6a0b78925ddb13b4dd1c87d9c5f1cbb098843d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Praszmo?= <michalpr@cert.pl>
Date: Tue, 24 May 2022 16:57:43 +0200
Subject: [PATCH] Add ECB mode to DES3 (#76)

---
 malduck/crypto/des3.py | 37 +++++++++++++++++++++++++++++++++++++
 tests/test_crypto.py   | 18 ++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/malduck/crypto/des3.py b/malduck/crypto/des3.py
index b61fb6d..8f161a5 100644
--- a/malduck/crypto/des3.py
+++ b/malduck/crypto/des3.py
@@ -6,6 +6,7 @@
 from Cryptodome.Cipher import DES
 from Cryptodome.Cipher import DES3 as DES3Cipher
 from Cryptodome.Cipher._mode_cbc import CbcMode
+from Cryptodome.Cipher._mode_ecb import EcbMode
 
 __all__ = ["des3"]
 
@@ -49,8 +50,44 @@ def decrypt(self, key: bytes, iv: bytes, data: bytes) -> bytes:
         return self._get_cipher(key, iv).decrypt(data)
 
 
+class Des3Ecb:
+    def _get_cipher(self, key: bytes) -> EcbMode:
+        if len(key) == 8:
+            # For 8 bytes it fallbacks to single DES
+            # (original cryptography behaviour)
+            return cast(EcbMode, DES.new(key, DES.MODE_ECB))
+        return cast(EcbMode, DES3Cipher.new(key, DES3Cipher.MODE_ECB))
+
+    def encrypt(self, key: bytes, data: bytes) -> bytes:
+        """
+        Encrypts buffer using DES/DES3 algorithm in ECB mode.
+
+        :param key: Cryptographic key (16 or 24 bytes, 8 bytes for single DES)
+        :type key: bytes
+        :param data: Buffer to be encrypted
+        :type data: bytes
+        :return: Encrypted data
+        :rtype: bytes
+        """
+        return self._get_cipher(key).encrypt(data)
+
+    def decrypt(self, key: bytes, data: bytes) -> bytes:
+        """
+        Decrypts buffer using DES/DES3 algorithm in ECB mode.
+
+        :param key: Cryptographic key (16 or 24 bytes, 8 bytes for single DES)
+        :type key: bytes
+        :param data: Buffer to be decrypted
+        :type data: bytes
+        :return: Decrypted data
+        :rtype: bytes
+        """
+        return self._get_cipher(key).decrypt(data)
+
+
 class Des3:
     cbc = Des3Cbc()
+    ecb = Des3Ecb()
 
 
 des3 = Des3()
diff --git a/tests/test_crypto.py b/tests/test_crypto.py
index 0cc6c48..dcc028e 100644
--- a/tests/test_crypto.py
+++ b/tests/test_crypto.py
@@ -170,6 +170,24 @@ def test_des():
     ) == b"\x1d\xed\xc37pV\x89S\xac\xaeT\xaf\xa1\xcfW\xa3"
 
 
+def test_des3():
+    assert des3.cbc.decrypt(
+        b"1"*8+b"3"*8+b"5"*8, b"B"*8, b'\xca\x8fVk\x11\xed"\x9c\xe7^T\x1c\xce\xae`\xee'
+    ) == b"C"*16
+
+    assert des3.cbc.encrypt(
+        b"1"*8+b"3"*8+b"5"*8,  b"B"*8, b"C"*16
+    ) == b'\xca\x8fVk\x11\xed"\x9c\xe7^T\x1c\xce\xae`\xee'
+
+    assert des3.ecb.decrypt(
+        b"1"*8+b"3"*8+b"5"*8, b'\xd4\x12\x80\xbaW\xd8g\xee\xd4\x12\x80\xbaW\xd8g\xee'
+    ) == b"C"*16
+
+    assert des3.ecb.encrypt(
+        b"1"*8+b"3"*8+b"5"*8, b"C"*16
+    ) == b'\xd4\x12\x80\xbaW\xd8g\xee\xd4\x12\x80\xbaW\xd8g\xee'
+
+
 def test_chacha20():
     assert chacha20.decrypt(
         key=b"A"*32, data=b'P\xb6\x12W\xf4\xd7\x83|,\xea\x04n\xba\x08Kj', nonce=b"C"*8