Add Encryptor class

sdk/python/aleo/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
 from ._aleolib import *
+from .encryptor import *
 
 __doc__ = _aleolib.__doc__
-

sdk/python/aleo/encryptor.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from ._aleolib import PrivateKey, Ciphertext, Field, Network, Identifier, Plaintext, Literal
+
+
+class Encryptor:
+    @staticmethod
+    # Encrypt a private key into ciphertext using a secret
+    def encrypt_private_key_with_secret(private_key: PrivateKey, secret: str) -> Ciphertext:
+        seed = private_key.seed()
+        return Encryptor.__encrypt_field(seed, secret, "private_key")
+
+    @staticmethod
+    # Decrypt a private key from ciphertext using a secret
+    def decrypt_private_key_with_secret(ciphertext: Ciphertext, secret: str) -> PrivateKey:
+        seed = Encryptor.__decrypt_field(ciphertext, secret, "private_key")
+        return PrivateKey.from_seed(seed)
+
+    @staticmethod
+    # Encrypted a field element into a ciphertext representation
+    def __encrypt_field(field: Field, secret: str, domain: str) -> Ciphertext:
+        domain_f = Field.domain_separator(domain)
+        secret_f = Field.domain_separator(secret)
+
+        nonce = Field.random()
+        blinding = Network.hash_psd2([domain_f, nonce, secret_f])
+        key = blinding * field
+        key_kv = (Identifier.from_string("key"),
+                  Plaintext.new_literal(Literal.from_field(key)))
+        nonce_kv = (Identifier.from_string("nonce"),
+                  Plaintext.new_literal(Literal.from_field(nonce)))
+        plaintext = Plaintext.new_struct([key_kv, nonce_kv])
+        return plaintext.encrypt_symmetric(secret_f)
+
+    @staticmethod
+    def __extract_value(plaintext: Plaintext, identifier: str) -> Field:
+        assert plaintext.is_struct()
+        ident = Identifier.from_string(identifier)
+        dec_map = plaintext.as_struct()
+        val = dec_map[ident]
+        assert val.is_literal()
+        literal = val.as_literal()
+        assert literal.type_name() == 'field'
+        return Field.from_string(str(literal))
+
+    @staticmethod
+    # Recover a field element encrypted within ciphertext
+    def __decrypt_field(ciphertext: Ciphertext, secret: str, domain: str) -> Field:
+        domain_f = Field.domain_separator(domain)
+        secret_f = Field.domain_separator(secret)
+        decrypted = ciphertext.decrypt_symmetric(secret_f)
+        assert decrypted.is_struct()
+        recovered_key = Encryptor.__extract_value(decrypted, "key")
+        recovered_nonce = Encryptor.__extract_value(decrypted, "nonce")
+        recovered_blinding = Network.hash_psd2([domain_f, recovered_nonce, secret_f])
+        return recovered_key / recovered_blinding
+

sdk/python/test.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
-import aleo
 import unittest
-
+import aleo
 
 class TestAleo(unittest.TestCase):
 
@@ -66,6 +65,20 @@ def test_account_sanity(self):
         self.assertFalse(account.verify(signature, bad_message))
         self.assertTrue(signature.verify(account.address(), message))
 
+    def test_encrypt_decrypt_sk(self):
+        private_key = aleo.PrivateKey.from_string(
+            "APrivateKey1zkpJYx2NZeJYB74JHpzvQGpKneTP75Dk8dao6paugZXtCz3")
+        ciphertext = aleo.Ciphertext.from_string(
+            "ciphertext1qvqt0sp0pp49gjeh50alfalt7ug3g8y7ha6cl3jkavcsnz8d0y9jwr27taxfrwd5kly8lah53qure3vxav6zxr7txattdvscv0kf3vcuqv9cmzj32znx4uwxdawcj3273zhgm8qwpxqczlctuvjvc596mgsqjxwz37f")
+        recovered = aleo.Encryptor.decrypt_private_key_with_secret(ciphertext, "qwe123")
+
+        self.assertEqual(private_key, recovered)
+
+        encrypted = aleo.Encryptor.encrypt_private_key_with_secret(private_key, "asd123")
+        other_recovered = aleo.Encryptor.decrypt_private_key_with_secret(encrypted, "asd123")
+
+        self.assertEqual(private_key, other_recovered)
+
     def test_coinbase(self):
         address = aleo.Address.from_string(
             "aleo16xwtrvntrfnan84sy3qg2gdkkp5u5p7sjc882lx8n06fjx2k0yqsklw8sv")