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

Fixes #124

3 files changed, 72 insertions, 2 deletions

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,