Add functionality to extract EC key from public key + Update tests

Fixes #124

5 files changed, 136 insertions, 9 deletions

diff --git a/MANIFEST.in b/MANIFEST.in
index 1eb138a..e14b231 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,6 +1,6 @@
-include SWIG/*.i
-include SWIG/*.h
-include SWIG/*.def
+include src/SWIG/*.i
+include src/SWIG/*.h
+include src/SWIG/*.def
 recursive-include tests *.py *.pem *.der *.b64 README *.pgp *.dat *.p7* *.crt *.txt
 recursive-include doc *
 recursive-include contrib *
@@ -9,5 +9,5 @@ include README.rst
 include CHANGES
 include epydoc.conf
 include LICENCE
-include SWIG/_m2crypto_wrap.c
-include M2Crypto/m2crypto.py
+include src/SWIG/_m2crypto_wrap.c
+include src/M2Crypto/m2crypto.py
diff --git a/src/M2Crypto/EC.py b/src/M2Crypto/EC.py
index 522092a..8cd052a 100644
--- a/src/M2Crypto/EC.py
+++ b/src/M2Crypto/EC.py
@@ -12,7 +12,6 @@ All rights reserved."""
 
 from M2Crypto import BIO, Err, EVP, m2, util
 from typing import AnyStr, Callable, Dict, Optional, Tuple, Union  # noqa
-from M2Crypto.EVP import PKey
 
 EC_Key = bytes
 
@@ -322,6 +321,16 @@ class EC_pub(EC):
         assert self.check_key(), 'key is not initialised'
         return m2.ec_key_get_public_key(self.ec)
 
+    def as_pem(self):
+        """
+        Returns the key(pair) as a string in PEM format.
+        If no password is passed and the cipher is set
+        it exits with error
+        """
+        with BIO.MemoryBuffer() as bio:
+            self.save_key_bio(bio)
+            return bio.read()
+
     save_key = EC.save_pub_key
 
     save_key_bio = EC.save_pub_key_bio
diff --git a/src/M2Crypto/EVP.py b/src/M2Crypto/EVP.py
index c48b670..de2dcef 100644
--- a/src/M2Crypto/EVP.py
+++ b/src/M2Crypto/EVP.py
@@ -9,7 +9,7 @@ Author: Heikki Toivonen
 """
 
 import logging
-from M2Crypto import BIO, Err, RSA, m2, util
+from M2Crypto import BIO, Err, RSA, EC, m2, util
 from typing import AnyStr, Optional, Callable  # noqa
 
 log = logging.getLogger('EVP')
@@ -385,6 +385,38 @@ class PKey(object):
         rsa = RSA.RSA_pub(rsa_ptr, 1)
         return rsa
 
+    def assign_ec(self, ec, capture=1):
+        # type: (EC.EC, int) -> int
+        """
+        Assign the EC key pair to self.
+
+        :param ec: M2Crypto.EC.EC object to be assigned to self.
+
+        :param capture: If true (default), this PKey object will own the EC
+                        object, meaning that once the PKey object gets
+                        deleted it is no longer safe to use the EC object.
+
+        :return: Return 1 for success and 0 for failure.
+        """
+        if capture:
+            ret = m2.pkey_assign_ec(self.pkey, ec.ec)
+            if ret:
+                ec._pyfree = 0
+        else:
+            ret = m2.pkey_set1_ec(self.pkey, ec.ec)
+        return ret
+
+    def get_ec(self):
+        # type: () -> EC.EC_pub
+        """
+        Return the underlying EC key if that is what the EVP
+        instance is holding.
+        """
+        ec_ptr = m2.pkey_get1_ec(self.pkey)
+
+        ec = EC.EC_pub(ec_ptr)
+        return ec
+
     def save_key(self, file, cipher='aes_128_cbc',
                  callback=util.passphrase_callback):
         # type: (AnyStr, Optional[str], Callable) -> int
diff --git a/src/SWIG/_evp.i b/src/SWIG/_evp.i
index 758f11b..cb480f3 100644
--- a/src/SWIG/_evp.i
+++ b/src/SWIG/_evp.i
@@ -17,6 +17,7 @@ Copyright (c) 2009-2010 Heikki Toivonen. All rights reserved.
 #include <openssl/evp.h>
 #include <openssl/hmac.h>
 #include <openssl/rsa.h>
+#include <openssl/ec.h>
 #include <openssl/opensslv.h>
 
 #if OPENSSL_VERSION_NUMBER < 0x10100000L
@@ -51,6 +52,7 @@ typedef struct evp_md_ctx_st EVP_MD_CTX;
 %apply Pointer NONNULL { EVP_CIPHER_CTX * };
 %apply Pointer NONNULL { EVP_CIPHER * };
 %apply Pointer NONNULL { RSA * };
+%apply Pointer NONNULL { EC_KEY * };
 
 %rename(md5) EVP_md5;
 extern const EVP_MD *EVP_md5(void);
@@ -177,6 +179,8 @@ extern int EVP_PKEY_assign(EVP_PKEY *, int, char *);
 #if OPENSSL_VERSION_NUMBER >= 0x0090800fL && !defined(OPENSSL_NO_EC)
 %rename(pkey_assign_ec) EVP_PKEY_assign_EC_KEY;
 extern int EVP_PKEY_assign_EC_KEY(EVP_PKEY *, EC_KEY *);
+%rename(pkey_set1_ec) EVP_PKEY_set1_EC_KEY;
+extern int EVP_PKEY_set1_EC_KEY(EVP_PKEY *, EC_KEY *);
 #endif
 %rename(pkey_set1_rsa) EVP_PKEY_set1_RSA;
 extern int EVP_PKEY_set1_RSA(EVP_PKEY *, RSA *);
@@ -228,6 +232,31 @@ RSA *pkey_get1_rsa(EVP_PKEY *pkey) {
 %}
 %typemap(out) RSA * ;
 
+%typemap(out) EC_KEY * {
+    PyObject *self = NULL; /* bug in SWIG_NewPointerObj as of 3.0.5 */
+
+    if ($1 != NULL)
+        $result = SWIG_NewPointerObj($1, $1_descriptor, 0);
+    else {
+        $result = NULL;
+    }
+}
+%inline %{
+EC_KEY *pkey_get1_ec(EVP_PKEY *pkey) {
+    EC_KEY *ret = NULL;
+
+    if ((ret = EVP_PKEY_get1_EC_KEY(pkey)) == NULL) {
+        /* _evp_err now inherits from PyExc_ValueError, so we should
+         * keep API intact.
+         */
+        PyErr_Format(_evp_err, "Invalid key in function %s.", __FUNCTION__);
+    }
+
+    return ret;
+}
+%}
+%typemap(out) EC_KEY * ;
+
 %inline %{
 PyObject *pkcs5_pbkdf2_hmac_sha1(PyObject *pass,
                                  PyObject *salt,
diff --git a/tests/test_evp.py b/tests/test_evp.py
index a355af7..75ed1b8 100644
--- a/tests/test_evp.py
+++ b/tests/test_evp.py
@@ -14,7 +14,7 @@ import logging
 
 from binascii import a2b_hex, hexlify, unhexlify
 
-from M2Crypto import BIO, EVP, RSA, Rand, m2, six, util
+from M2Crypto import BIO, EVP, RSA, EC, Rand, m2, six, util
 from tests import unittest
 from tests.fips import fips_mode
 
@@ -196,7 +196,7 @@ class EVPTestCase(unittest.TestCase):
         rsa3 = RSA.gen_key(1024, 3, callback=self._gen_callback)
         self.assertNotEqual(rsa.sign(digest), rsa3.sign(digest))
 
-    def test_load_key_string_pubkey(self):
+    def test_load_key_string_pubkey_rsa(self):
         """
         Testing creating a PKey instance from PEM string.
         """
@@ -225,6 +225,63 @@ class EVPTestCase(unittest.TestCase):
         with self.assertRaises(ValueError):
             pkey.get_rsa()
 
+    def test_get_ec(self):
+        """
+        Testing retrieving the EC key from the PKey instance.
+        """
+        ec = EC.gen_params(m2.NID_secp112r1)
+        ec.gen_key()
+        self.assertIsInstance(ec, EC.EC)
+        pkey = EVP.PKey()
+        pkey.assign_ec(ec)
+        ec2 = pkey.get_ec()
+        self.assertIsInstance(ec2, EC.EC_pub)
+        self.assertEqual(ec.compute_dh_key(ec), ec2.compute_dh_key(ec2))
+
+        pem = ec.as_pem(callback=self._pass_callback)
+        pem2 = ec2.as_pem()
+        assert pem
+        assert pem2
+        self.assertNotEqual(pem, pem2)
+
+        message = b'This is the message string'
+        digest = hashlib.sha1(message).digest()
+        ec_sign = ec.sign_dsa(digest)
+        ec2_sign = ec.sign_dsa(digest)
+        self.assertEqual(ec.verify_dsa(digest, ec_sign[0], ec_sign[1]), ec.verify_dsa(digest, ec2_sign[0], ec2_sign[1]))
+
+        ec3 = EC.gen_params(m2.NID_secp112r1)
+        ec3.gen_key()
+        ec3_sign = ec.sign_dsa(digest)
+        self.assertEqual(ec.verify_dsa(digest, ec_sign[0], ec_sign[1]), ec.verify_dsa(digest, ec3_sign[0], ec3_sign[1]))
+
+    def test_load_key_string_pubkey_ec(self):
+        """
+        Testing creating a PKey instance from PEM string.
+        """
+        ec = EC.gen_params(m2.NID_secp112r1)
+        ec.gen_key()
+        self.assertIsInstance(ec, EC.EC)
+        ec_pem = BIO.MemoryBuffer()
+        ec.save_pub_key_bio(ec_pem)
+        pkey = EVP.load_key_string_pubkey(ec_pem.read())
+        ec2 = pkey.get_ec()
+        self.assertIsInstance(ec2, EC.EC_pub)
+        pem = ec.as_pem(callback=self._pass_callback)
+        pem2 = ec2.as_pem()
+        assert pem
+        assert pem2
+        self.assertNotEqual(pem, pem2)
+
+    def test_get_ec_fail(self):
+        """
+        Testing trying to retrieve the EC key from the PKey instance
+        when it is not holding a EC Key. Should raise a ValueError.
+        """
+        pkey = EVP.PKey()
+        with self.assertRaises(ValueError):
+            pkey.get_ec()
+
     def test_get_modulus(self):
         pkey = EVP.PKey()
         with self.assertRaises(ValueError):