From be45b18258534af40bc8c80dda99fd5fb77c8bea Mon Sep 17 00:00:00 2001
From: tuxxy <me@johnpacific.com>
Date: Tue, 10 Oct 2017 17:20:48 -0600
Subject: [PATCH] Add tests for ecdsa_sign and ecdsa_verify

---
 nkms/crypto/crypto.py       | 11 +++++++----
 tests/crypto/test_crypto.py | 24 ++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/nkms/crypto/crypto.py b/nkms/crypto/crypto.py
index 7b467088c..6314838f6 100644
--- a/nkms/crypto/crypto.py
+++ b/nkms/crypto/crypto.py
@@ -169,15 +169,19 @@ def ecdsa_sign(
 
 
 def ecdsa_verify(
-    signature: bytes,
+    v: int,
+    r: int,
+    s: int,
     msghash: bytes,
     pubkey: Union[bytes, Tuple[int]]
 ) -> bool:
     """
     Takes a msgpacked signature and verifies the message.
 
+    :param v: V of sig
+    :param r: R of sig
+    :param s: S of sig
     :param bytes msghash: The hashed message to verify
-    :param bytes signature: The msgpacked signature (v, r, and s)
     :param bytes pubkey: Pubkey to validate signature for
 
     :rtype: Boolean
@@ -186,8 +190,7 @@ def ecdsa_verify(
     if bytes == type(pubkey):
         pubkey = ecdsa_bytes2pub(pubkey)
 
-    sig = ecdsa_load_sig(signature)
-    verify_sig = ecdsa_raw_recover(msghash, sig)
+    verify_sig = ecdsa_raw_recover(msghash, (v, r, s))
     # TODO: Should this equality test be done better?
     return verify_sig == pubkey
 
diff --git a/tests/crypto/test_crypto.py b/tests/crypto/test_crypto.py
index 4128d75bd..448670c4d 100644
--- a/tests/crypto/test_crypto.py
+++ b/tests/crypto/test_crypto.py
@@ -115,9 +115,29 @@ class TestCrypto(unittest.TestCase):
         self.assertEqual(3, len(loaded_sig))
         self.assertEqual((1, 2, 3), loaded_sig)
 
+    def test_ecdsa_sign(self):
+        msghash = Crypto.secure_random(32)
+        privkey = Crypto.ecdsa_gen_priv()
+
+        vrs = Crypto.ecdsa_sign(msghash, privkey)
+        self.assertEqual(tuple, type(vrs))
+        self.assertEqual(3, len(vrs))
+
+    def test_ecdsa_verify(self):
+        msghash = Crypto.secure_random(32)
+        privkey = Crypto.ecdsa_gen_priv()
+        pubkey = Crypto.ecdsa_priv2pub(privkey, to_bytes=False)
+
+        vrs = Crypto.ecdsa_sign(msghash, privkey)
+        self.assertEqual(tuple, type(vrs))
+        self.assertEqual(3, len(vrs))
+
+        is_verified = Crypto.ecdsa_verify(*vrs, msghash, pubkey)
+        self.assertEqual(bool, type(is_verified))
+        self.assertTrue(is_verified)
 
     def test_symm_encrypt(self):
-        key = random._urandom(32)
+        key = Crypto.secure_random(32)
         plaintext = b'this is a test'
 
         ciphertext = Crypto.symm_encrypt(key, plaintext)
@@ -125,7 +145,7 @@ class TestCrypto(unittest.TestCase):
         self.assertNotEqual(plaintext, ciphertext)
 
     def test_symm_decrypt(self):
-        key = random._urandom(32)
+        key = Crypto.secure_random(32)
         plaintext = b'this is a test'
 
         ciphertext = Crypto.symm_encrypt(key, plaintext)