[PrivateKey] Add support for AIP80 to Ed25519 and Secp256k1

aptos_sdk/asymmetric_crypto.py

@@ -3,18 +3,108 @@
 
 from __future__ import annotations
 
+from enum import Enum
+
 from typing_extensions import Protocol
 
 from .bcs import Deserializable, Serializable
 
 
+class PrivateKeyVariant(Enum):
+    Ed25519 = "ed25519"
+    Secp256k1 = "secp256k1"
+
+
 class PrivateKey(Deserializable, Serializable, Protocol):
     def hex(self) -> str: ...
 
     def public_key(self) -> PublicKey: ...
 
     def sign(self, data: bytes) -> Signature: ...
 
+    """
+    The AIP-80 compliant prefixes for each private key type. Append this to a private key's hex representation
+    to get an AIP-80 compliant string.
+
+    [Read about AIP-80](https://github.com/aptos-foundation/AIPs/blob/main/aips/aip-80.md)
+    """
+    AIP80_PREFIXES: dict[PrivateKeyVariant, str] = {
+        PrivateKeyVariant.Ed25519: "ed25519-priv-",
+        PrivateKeyVariant.Secp256k1: "secp256k1-priv-",
+    }
+
+    @staticmethod
+    def format_private_key(
+        private_key: bytes | str, key_type: PrivateKeyVariant
+    ) -> str:
+        """
+        Format a HexInput to an AIP-80 compliant string.
+
+        :param private_key: The hex string or bytes format of the private key.
+        :param key_type: The private key type.
+        :return: AIP-80 compliant string.
+        """
+        if key_type not in PrivateKey.AIP80_PREFIXES:
+            raise ValueError(f"Unknown private key type: {key_type}")
+        aip80_prefix = PrivateKey.AIP80_PREFIXES[key_type]
+
+        key_value: str | None = None
+        if isinstance(private_key, str):
+            key_value = private_key
+        elif isinstance(private_key, bytes):
+            key_value = f"0x{private_key.hex()}"
+        else:
+            raise TypeError("Input value must be a string or bytes.")
+
+        return f"{aip80_prefix}{key_value}"
+
+    @staticmethod
+    def parse_hex_input(
+        value: str | bytes, key_type: PrivateKeyVariant, strict: bool | None = None
+    ) -> bytes:
+        """
+        Parse a HexInput that may be a hex string, bytes, or an AIP-80 compliant string to a byte array.
+
+        :param value: A hex string, byte array, or AIP-80 compliant string.
+        :param key_type: The private key type.
+        :param strict: If true, the value MUST be compliant with AIP-80.
+        :return: Parsed private key as bytes.
+        """
+        if key_type not in PrivateKey.AIP80_PREFIXES:
+            raise ValueError(f"Unknown private key type: {key_type}")
+        aip80_prefix = PrivateKey.AIP80_PREFIXES[key_type]
+
+        if isinstance(value, str):
+            if not strict and not value.startswith(aip80_prefix):
+                # Non-AIP-80 compliant hex string
+                if strict is None:
+                    print(
+                        "It is recommended that private keys are AIP-80 compliant (https://github.com/aptos-foundation/AIPs/blob/main/aips/aip-80.md)."
+                    )
+                if value[0:2] == "0x":
+                    value = value[2:]
+                return bytes.fromhex(value)
+            elif value.startswith(aip80_prefix):
+                # AIP-80 compliant string
+                value = value.split("-")[2]
+                if value[0:2] == "0x":
+                    value = value[2:]
+                return bytes.fromhex(value)
+            else:
+                if strict:
+                    raise ValueError(
+                        "Invalid HexString input. Must be AIP-80 compliant string."
+                    )
+                raise ValueError("Invalid HexString input.")
+        elif isinstance(value, bytes):
+            if strict is None:
+                print(
+                    "It is recommended that private keys are AIP-80 compliant (https://github.com/aptos-foundation/AIPs/blob/main/aips/aip-80.md)."
+                )
+            return value
+        else:
+            raise TypeError("Input value must be a string or bytes.")
+
 
 class PublicKey(Deserializable, Serializable, Protocol):
     def to_crypto_bytes(self) -> bytes:

aptos_sdk/ed25519.py

@@ -29,14 +29,41 @@ def __str__(self):
         return self.hex()
 
     @staticmethod
-    def from_str(value: str) -> PrivateKey:
-        if value[0:2] == "0x":
-            value = value[2:]
-        return PrivateKey(SigningKey(bytes.fromhex(value)))
+    def from_hex(value: str | bytes, strict: bool | None = None) -> PrivateKey:
+        """
+        Parse a HexInput that may be a hex string, bytes, or an AIP-80 compliant string to a private key.
+
+        :param value: A hex string, byte array, or AIP-80 compliant string.
+        :param strict: If true, the value MUST be compliant with AIP-80.
+        :return: Parsed Ed25519 private key.
+        """
+        return PrivateKey(
+            SigningKey(
+                PrivateKey.parse_hex_input(
+                    value, asymmetric_crypto.PrivateKeyVariant.Ed25519, strict
+                )
+            )
+        )
+
+    @staticmethod
+    def from_str(value: str, strict: bool | None = None) -> PrivateKey:
+        """
+        Parse a HexInput that may be a hex string or an AIP-80 compliant string to a private key.
+
+        :param value: A hex string or AIP-80 compliant string.
+        :param strict: If true, the value MUST be compliant with AIP-80.
+        :return: Parsed Ed25519 private key.
+        """
+        return PrivateKey.from_hex(value, strict)
 
     def hex(self) -> str:
         return f"0x{self.key.encode().hex()}"
 
+    def aip80(self) -> str:
+        return PrivateKey.format_private_key(
+            self.hex(), asymmetric_crypto.PrivateKeyVariant.Ed25519
+        )
+
     def public_key(self) -> PublicKey:
         return PublicKey(self.key.verify_key)
 
@@ -278,6 +305,26 @@ def serialize(self, serializer: Serializer):
 
 
 class Test(unittest.TestCase):
+    def test_private_key_from_str(self):
+        private_key_hex = PrivateKey.from_str(
+            "0x4e5e3be60f4bbd5e98d086d932f3ce779ff4b58da99bf9e5241ae1212a29e5fe", False
+        )
+        private_key_with_prefix = PrivateKey.from_str(
+            "ed25519-priv-0x4e5e3be60f4bbd5e98d086d932f3ce779ff4b58da99bf9e5241ae1212a29e5fe",
+            True,
+        )
+        private_key_bytes = PrivateKey.from_hex(
+            bytes.fromhex(
+                "4e5e3be60f4bbd5e98d086d932f3ce779ff4b58da99bf9e5241ae1212a29e5fe"
+            ),
+            False,
+        )
+        self.assertEqual(
+            private_key_hex.hex(),
+            private_key_with_prefix.hex(),
+            private_key_bytes.hex(),
+        )
+
     def test_sign_and_verify(self):
         in_value = b"test_message"
 
@@ -317,10 +364,10 @@ def test_signature_key_serialization(self):
     def test_multisig(self):
         # Generate signatory private keys.
         private_key_1 = PrivateKey.from_str(
-            "4e5e3be60f4bbd5e98d086d932f3ce779ff4b58da99bf9e5241ae1212a29e5fe"
+            "ed25519-priv-0x4e5e3be60f4bbd5e98d086d932f3ce779ff4b58da99bf9e5241ae1212a29e5fe"
         )
         private_key_2 = PrivateKey.from_str(
-            "1e70e49b78f976644e2c51754a2f049d3ff041869c669523ba95b172c7329901"
+            "ed25519-priv-0x1e70e49b78f976644e2c51754a2f049d3ff041869c669523ba95b172c7329901"
         )
         # Generate multisig public key with threshold of 1.
         multisig_public_key = MultiPublicKey(

aptos_sdk/secp256k1_ecdsa.py

@@ -30,18 +30,42 @@ def __str__(self):
         return self.hex()
 
     @staticmethod
-    def from_str(value: str) -> PrivateKey:
-        if value[0:2] == "0x":
-            value = value[2:]
-        if len(value) != PrivateKey.LENGTH * 2:
+    def from_hex(value: str | bytes, strict: bool | None = None) -> PrivateKey:
+        """
+        Parse a HexInput that may be a hex string, bytes, or an AIP-80 compliant string to a private key.
+
+        :param value: A hex string, byte array, or AIP-80 compliant string.
+        :param strict: If true, the value MUST be compliant with AIP-80.
+        :return: Parsed private key as bytes.
+        """
+        parsed_value = PrivateKey.parse_hex_input(
+            value, asymmetric_crypto.PrivateKeyVariant.Secp256k1, strict
+        )
+        if len(parsed_value.hex()) != PrivateKey.LENGTH * 2:
             raise Exception("Length mismatch")
         return PrivateKey(
-            SigningKey.from_string(bytes.fromhex(value), SECP256k1, hashlib.sha3_256)
+            SigningKey.from_string(parsed_value, SECP256k1, hashlib.sha3_256)
         )
 
+    @staticmethod
+    def from_str(value: str, strict: bool | None = None) -> PrivateKey:
+        """
+        Parse a HexInput that may be a hex string or an AIP-80 compliant string to a private key.
+
+        :param value: A hex string or AIP-80 compliant string.
+        :param strict: If true, the value MUST be compliant with AIP-80.
+        :return: Parsed Secp256k1 private key.
+        """
+        return PrivateKey.from_hex(value, strict)
+
     def hex(self) -> str:
         return f"0x{self.key.to_string().hex()}"
 
+    def aip80(self) -> str:
+        return PrivateKey.format_private_key(
+            self.hex(), asymmetric_crypto.PrivateKeyVariant.Ed25519
+        )
+
     def public_key(self) -> PublicKey:
         return PublicKey(self.key.verifying_key)
 
@@ -177,10 +201,28 @@ def serialize(self, serializer: Serializer):
 
 
 class Test(unittest.TestCase):
-    def test_vectors(self):
-        private_key_hex = (
-            "0x306fa009600e27c09d2659145ce1785249360dd5fb992da01a578fe67ed607f4"
+    def test_private_key_from_str(self):
+        private_key_hex = PrivateKey.from_str(
+            "0x306fa009600e27c09d2659145ce1785249360dd5fb992da01a578fe67ed607f4", False
+        )
+        private_key_with_prefix = PrivateKey.from_str(
+            "secp256k1-priv-0x306fa009600e27c09d2659145ce1785249360dd5fb992da01a578fe67ed607f4",
+            True,
+        )
+        private_key_bytes = PrivateKey.from_hex(
+            bytes.fromhex(
+                "306fa009600e27c09d2659145ce1785249360dd5fb992da01a578fe67ed607f4"
+            ),
+            False,
         )
+        self.assertEqual(
+            private_key_hex.hex(),
+            private_key_with_prefix.hex(),
+            private_key_bytes.hex(),
+        )
+
+    def test_vectors(self):
+        private_key_hex = "secp256k1-priv-0x306fa009600e27c09d2659145ce1785249360dd5fb992da01a578fe67ed607f4"
         public_key_hex = "0x04210c9129e35337ff5d6488f90f18d842cf985f06e0baeff8df4bfb2ac4221863e2631b971a237b5db0aa71188e33250732dd461d56ee623cbe0426a5c2db79ef"
         signature_hex = "0xa539b0973e76fa99b2a864eebd5da950b4dfb399c7afe57ddb34130e454fc9db04dceb2c3d4260b8cc3d3952ab21b5d36c7dc76277fe3747764e6762d12bd9a9"
         data = b"Hello world"