From 623af875ef67e7886e9acb35d5efc03ca8b178d4 Mon Sep 17 00:00:00 2001
From: Alexandre Paillier <alexandre.paillier@ledger.fr>
Date: Wed, 30 Nov 2022 11:39:06 +0100
Subject: [PATCH] Ragger tests can now handle multiple signing keys

---
 tests/ragger/cal/cal.py                    | 18 ------------------
 tests/ragger/eip712/InputData.py           |  6 +++---
 tests/ragger/keys.py                       | 22 ++++++++++++++++++++++
 tests/ragger/{cal/key.pem => keys/cal.pem} |  0
 4 files changed, 25 insertions(+), 21 deletions(-)
 delete mode 100644 tests/ragger/cal/cal.py
 create mode 100644 tests/ragger/keys.py
 rename tests/ragger/{cal/key.pem => keys/cal.pem} (100%)

diff --git a/tests/ragger/cal/cal.py b/tests/ragger/cal/cal.py
deleted file mode 100644
index e75fcabd..00000000
--- a/tests/ragger/cal/cal.py
+++ /dev/null
@@ -1,18 +0,0 @@
-import os
-import hashlib
-from ecdsa.util import sigencode_der
-from ecdsa import SigningKey
-
-_key: SigningKey = None
-
-def _init_key():
-    global _key
-    with open(os.path.dirname(__file__) + "/key.pem") as pem_file:
-        _key = SigningKey.from_pem(pem_file.read(), hashlib.sha256)
-    assert _key != None
-
-def sign(data: bytes) -> bytes:
-    global _key
-    if not _key:
-        _init_key()
-    return _key.sign_deterministic(data, sigencode=sigencode_der)
diff --git a/tests/ragger/eip712/InputData.py b/tests/ragger/eip712/InputData.py
index 40a38825..ae44bc62 100644
--- a/tests/ragger/eip712/InputData.py
+++ b/tests/ragger/eip712/InputData.py
@@ -5,7 +5,7 @@ import sys
 import re
 import hashlib
 from app.client import EthereumClient, EIP712FieldType
-from cal import cal
+from keys import key_sign, Key
 
 # global variables
 app_client: EthereumClient = None
@@ -251,7 +251,7 @@ def send_filtering_message_info(display_name: str, filters_count: int):
     for char in display_name:
         to_sign.append(ord(char))
 
-    sig = cal.sign(to_sign)
+    sig = key_sign(Key.CAL, to_sign)
     app_client.eip712_filtering_message_info(display_name, filters_count, sig)
 
 # ledgerjs doesn't actually sign anything, and instead uses already pre-computed signatures
@@ -269,7 +269,7 @@ def send_filtering_show_field(display_name):
         to_sign.append(ord(char))
     for char in display_name:
         to_sign.append(ord(char))
-    sig = cal.sign(to_sign)
+    sig = key_sign(Key.CAL, to_sign)
     app_client.eip712_filtering_show_field(display_name, sig)
 
 def read_filtering_file(domain, message, filtering_file_path):
diff --git a/tests/ragger/keys.py b/tests/ragger/keys.py
new file mode 100644
index 00000000..f6ff5f2f
--- /dev/null
+++ b/tests/ragger/keys.py
@@ -0,0 +1,22 @@
+import os
+import hashlib
+from ecdsa.util import sigencode_der
+from ecdsa import SigningKey
+from enum import Enum, auto
+
+class Key(Enum):
+    CAL = auto()
+
+_keys: dict[Key, SigningKey] = dict()
+
+def _init_key(key: Key):
+    global _keys
+    with open("%s/keys/%s.pem" % (os.path.dirname(__file__), key.name.lower())) as pem_file:
+        _keys[key] = SigningKey.from_pem(pem_file.read(), hashlib.sha256)
+    assert (key in _keys) and (_keys[key] != None)
+
+def key_sign(key: Key, data: bytes) -> bytes:
+    global _keys
+    if key not in _keys:
+        _init_key(key)
+    return _keys[key].sign_deterministic(data, sigencode=sigencode_der)
diff --git a/tests/ragger/cal/key.pem b/tests/ragger/keys/cal.pem
similarity index 100%
rename from tests/ragger/cal/key.pem
rename to tests/ragger/keys/cal.pem