From dd9c8afdcb0d08b24f8f45f01270053031960ce6 Mon Sep 17 00:00:00 2001
From: jesko <rattle@nullteilerfrei.de>
Date: Tue, 10 Sep 2024 13:18:43 +0200
Subject: [PATCH] adds test and bugfix

---
 refinery/units/crypto/cipher/rc2.py  | 34 +++++++++++++++++++++++++++-
 test/units/crypto/cipher/test_rc2.py | 28 +++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 test/units/crypto/cipher/test_rc2.py

diff --git a/refinery/units/crypto/cipher/rc2.py b/refinery/units/crypto/cipher/rc2.py
index 76e5071f94..a88b6f25a0 100644
--- a/refinery/units/crypto/cipher/rc2.py
+++ b/refinery/units/crypto/cipher/rc2.py
@@ -3,6 +3,7 @@
 from Cryptodome.Cipher import ARC2
 
 from refinery.lib.crypto import PyCryptoFactoryWrapper
+from refinery.units import Arg
 from refinery.units.crypto.cipher import StandardBlockCipherUnit, CipherInterface
 
 
@@ -11,6 +12,37 @@ class rc2(StandardBlockCipherUnit, cipher=PyCryptoFactoryWrapper(ARC2)):
     RC2 encryption and decryption.
     """
 
+    def __init__(
+        self, key, iv=b'', *,
+        eks: Arg.Number('-k', '--eks', group='EKS',
+            help='Set the effective key size. Default is {default}.') = 1024,
+        derive_eks: Arg.Switch('-d', '--dks', group='EKS',
+            help='Act as .NET and derive the effective key size from the key length.') = False,
+        padding=None,
+        mode=None,
+        raw=False,
+        little_endian=False,
+        segment_size=0,
+        mac_len=0,
+        assoc_len=0,
+        **keywords
+    ):
+        super().__init__(
+            key,
+            iv,
+            eks=eks,
+            derive_eks=derive_eks,
+            padding=padding,
+            mode=mode,
+            raw=raw,
+            little_endian=little_endian,
+            segment_size=segment_size,
+            mac_len=mac_len,
+            assoc_len=assoc_len,
+            **keywords
+        )
+
     def _new_cipher(self, **optionals) -> CipherInterface:
-        optionals.update(effective_keylen=max(ARC2.key_size))
+        eks = len(self.args.key) * 8 if self.args.derive_eks else self.args.eks
+        optionals.update(effective_keylen=eks)
         return super()._new_cipher(**optionals)
diff --git a/test/units/crypto/cipher/test_rc2.py b/test/units/crypto/cipher/test_rc2.py
new file mode 100644
index 0000000000..529b8d5ed8
--- /dev/null
+++ b/test/units/crypto/cipher/test_rc2.py
@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+from ... import TestUnitBase
+
+
+class TestRC2(TestUnitBase):
+
+    def test_dotnet_eks_derivation(self):
+        data = b'U4dbuch3yV0yPCA7W+zIag'
+        rc2 = self.load(b'System', derive_eks=True, mode='cbc')
+        b64 = self.ldu('b64')
+        self.assertEqual(data | b64 | rc2 | str, 'InvokeMethod')
+
+    def test_bare_pycryptodome_vectors(self):
+        vectors = {
+            10: b'\xb8\x96\xb7\x2a\x5b\x4a\x03\x0c',
+            20: b'\x22\x0b\xfd\x7a\xf4\x32\x8c\xb0',
+            30: b'\xf8\x1d\xbe\xfe\xff\x9e\x70\xf4',
+            40: b'\x52\xe3\x32\xd1\xb5\xd6\xb0\x77',
+            50: b'\xb5\xc0\x92\xd6\xef\x75\x34\x85',
+            60: b'\xdc\xc3\xa7\xdd\xa6\x5f\x5b\x7a',
+            70: b'\x32\x55\x67\x68\x7a\x39\x8a\xfe',
+            80: b'\xac\xb0\xd2\xe4\x63\xbb\xec\xb2',
+            90: b'\x6a\xf7\x13\x91\xf0\xe1\x44\xbc',
+        }
+        for k, vec in vectors.items():
+            rc2 = self.load(B'A' * k, mode='ecb', reverse=True, raw=True)
+            self.assertEqual(bytes(range(8)) | rc2 | bytes, vec)