ecdsa key support

nadineloepfe · nadineloepfe · commit d2f2ce30ab3b · 2025-01-03T11:51:32.000+01:00
diff --git a/src/hedera_sdk_python/crypto/private_key.py b/src/hedera_sdk_python/crypto/private_key.py
@@ -1,50 +1,112 @@
-from cryptography.hazmat.primitives.asymmetric import ed25519
+from cryptography.hazmat.primitives.asymmetric import ed25519, ec
 from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.backends import default_backend
 from hedera_sdk_python.crypto.public_key import PublicKey
 
+
 class PrivateKey:
+    """
+    Represents a private key that can be either Ed25519 or ECDSA (secp256k1).
+    """
+
     def __init__(self, private_key):
+        """
+        Initializes a PrivateKey from a cryptography PrivateKey object.
+        """
         self._private_key = private_key
 
     @classmethod
-    def generate(cls):
+    def generate(cls, key_type: str = "ed25519"):
         """
-        Generates a new Ed25519 private key.
+        Generates a new private key. Defaults to an Ed25519 private key unless
+        'ecdsa' is specified.
+
+        Args:
+            key_type (str): Either 'ed25519' or 'ecdsa'. Defaults to 'ed25519'.
 
         Returns:
             PrivateKey: A new instance of PrivateKey.
         """
-        private_key = ed25519.Ed25519PrivateKey.generate()
+        if key_type.lower() == "ed25519":
+            return cls.generate_ed25519()
+        elif key_type.lower() == "ecdsa":
+            return cls.generate_ecdsa()
+        else:
+            raise ValueError("Invalid key_type. Use 'ed25519' or 'ecdsa'.")
+
+    @classmethod
+    def generate_ed25519(cls):
+        """
+        Generates a new Ed25519 private key.
+
+        Returns:
+            PrivateKey: A new instance of PrivateKey using Ed25519.
+        """
+        return cls(ed25519.Ed25519PrivateKey.generate())
+
+    @classmethod
+    def generate_ecdsa(cls):
+        """
+        Generates a new ECDSA (secp256k1) private key.
+
+        Returns:
+            PrivateKey: A new instance of PrivateKey using ECDSA.
+        """
+        private_key = ec.generate_private_key(ec.SECP256K1(), default_backend())
         return cls(private_key)
 
     @classmethod
     def from_string(cls, key_str):
         """
-        Load a private key from a hex-encoded string. Supports both raw private keys (32 bytes)
-        and DER-encoded private keys.
+        Load a private key from a hex-encoded string. For Ed25519, expects 32 bytes.
+        For ECDSA (secp256k1), also expects 32 bytes (raw scalar).
+        If it's DER-encoded, tries to parse and detect Ed25519 vs ECDSA.
+
+        Args:
+            key_str (str): The hex-encoded private key string.
+
+        Returns:
+            PrivateKey: A new instance of PrivateKey.
+
+        Raises:
+            ValueError: If the key is invalid or unsupported.
         """
         try:
             key_bytes = bytes.fromhex(key_str)
         except ValueError:
             raise ValueError("Invalid hex-encoded private key string.")
 
+        if len(key_bytes) == 32:
+            try:
+                ed_priv = ed25519.Ed25519PrivateKey.from_private_bytes(key_bytes)
+                return cls(ed_priv)
+            except Exception:
+                pass
+            try:
+                private_int = int.from_bytes(key_bytes, "big")
+                ec_priv = ec.derive_private_key(private_int, ec.SECP256K1(), default_backend())
+                return cls(ec_priv)
+            except Exception:
+                pass
+
         try:
-            if len(key_bytes) == 32:
-                private_key = ed25519.Ed25519PrivateKey.from_private_bytes(key_bytes)
-            else:
-                private_key = serialization.load_der_private_key(
-                    key_bytes, password=None
-                )
-                if not isinstance(private_key, ed25519.Ed25519PrivateKey):
-                    raise TypeError("The key is not an Ed25519 private key.")
-            return cls(private_key)
+            private_key = serialization.load_der_private_key(key_bytes, password=None)
         except Exception as e:
-            print(f"Error loading Ed25519 private key: {e}")
-            raise ValueError("Failed to load private key.")
+            raise ValueError(f"Failed to load private key (DER): {e}")
+
+        if isinstance(private_key, ed25519.Ed25519PrivateKey):
+            return cls(private_key)
+
+        if isinstance(private_key, ec.EllipticCurvePrivateKey):
+            if not isinstance(private_key.curve, ec.SECP256K1):
+                raise ValueError("Only secp256k1 ECDSA is supported.")
+            return cls(private_key)
+
+        raise ValueError("Unsupported private key type.")
 
-    def sign(self, data):
+    def sign(self, data: bytes) -> bytes:
         """
-        Signs the given data using the private key.
+        Signs the given data using this private key (Ed25519 or ECDSA).
 
         Args:
             data (bytes): The data to sign.
@@ -54,40 +116,63 @@ def sign(self, data):
         """
         return self._private_key.sign(data)
 
-    def public_key(self):
+    def public_key(self) -> PublicKey:
         """
-        Retrieves the corresponding public key.
+        Retrieves the corresponding PublicKey.
 
         Returns:
             PublicKey: The public key associated with this private key.
         """
         return PublicKey(self._private_key.public_key())
 
-    def to_string(self):
+    def to_bytes_raw(self) -> bytes:
         """
-        Returns the private key as a hex-encoded string.
+        Returns the private key bytes in raw form (32 bytes for both Ed25519 and ECDSA).
 
         Returns:
-            str: The hex-encoded private key.
+            bytes: The raw private key bytes.
         """
-        private_bytes = self._private_key.private_bytes(
+        return self._private_key.private_bytes(
             encoding=serialization.Encoding.Raw,
             format=serialization.PrivateFormat.Raw,
             encryption_algorithm=serialization.NoEncryption()
         )
-        return private_bytes.hex()
 
+    def to_string_raw(self) -> str:
+        """
+        Returns the raw private key as a hex-encoded string.
+
+        Returns:
+            str: The hex-encoded raw private key.
+        """
+        return self.to_bytes_raw().hex()
+    
+    def to_string(self) -> str:
+        """
+        Returns the private key as a hex string (raw).
+        """
+        return self.to_string_raw()
+
+    def is_ed25519(self) -> bool:
+        """
+        Checks if this private key is Ed25519.
+
+        Returns:
+            bool: True if Ed25519, False otherwise.
+        """
+        return isinstance(self._private_key, ed25519.Ed25519PrivateKey)
 
-    def to_bytes(self):
+    def is_ecdsa(self) -> bool:
         """
-        Returns the private key as bytes.
+        Checks if this private key is ECDSA (secp256k1).
 
         Returns:
-            bytes: The private key.
+            bool: True if ECDSA, False otherwise.
         """
-        private_bytes = self._private_key.private_bytes(
-            encoding=serialization.Encoding.Raw,
-            format=serialization.PrivateFormat.Raw,
-            encryption_algorithm=serialization.NoEncryption()
-        )
-        return private_bytes
+        from cryptography.hazmat.primitives.asymmetric import ec
+        return isinstance(self._private_key, ec.EllipticCurvePrivateKey)
+
+    def __repr__(self):
+        if self.is_ed25519():
+            return f"<PrivateKey (Ed25519) hex={self.to_string_raw()}>"
+        return f"<PrivateKey (ECDSA) hex={self.to_string_raw()}>"
diff --git a/src/hedera_sdk_python/crypto/public_key.py b/src/hedera_sdk_python/crypto/public_key.py
@@ -1,13 +1,21 @@
-from cryptography.hazmat.primitives.asymmetric import ed25519
+from cryptography.hazmat.primitives.asymmetric import ed25519, ec
 from cryptography.hazmat.primitives import serialization
 
 class PublicKey:
+    """
+    Represents a public key that can be either Ed25519 or ECDSA (secp256k1).
+    """
+
     def __init__(self, public_key):
+        """
+        Initializes a PublicKey from a cryptography PublicKey object.
+        """
         self._public_key = public_key
 
-    def verify(self, signature, data):
+    def verify(self, signature: bytes, data: bytes) -> None:
         """
-        Verifies a signature for the given data using the public key.
+        Verifies a signature for the given data using this public key.
+        Raises an exception if the signature is invalid.
 
         Args:
             signature (bytes): The signature to verify.
@@ -18,42 +26,73 @@ def verify(self, signature, data):
         """
         self._public_key.verify(signature, data)
 
-    def to_string(self):
+    def to_bytes_raw(self) -> bytes:
         """
-        Returns the public key as a hex-encoded string.
+        Returns the public key in raw form:
+         - For Ed25519, it's 32 bytes.
+         - For ECDSA (secp256k1), it's the uncompressed or compressed form, 
+           depending on how cryptography outputs RAW. Usually 33 bytes compressed.
 
         Returns:
-            str: The hex-encoded public key.
+            bytes: The raw public key bytes.
         """
-        public_bytes = self._public_key.public_bytes(
+        return self._public_key.public_bytes(
             encoding=serialization.Encoding.Raw,
             format=serialization.PublicFormat.Raw
         )
-        return public_bytes.hex()
 
-    def to_proto(self):
+    def to_string_raw(self) -> str:
         """
-        Returns the protobuf representation of the public key.
+        Returns the raw public key as a hex-encoded string.
 
         Returns:
-            Key: The protobuf Key message.
+            str: The hex-encoded raw public key.
         """
-        from hedera_sdk_python.hapi.services import basic_types_pb2
-        public_bytes = self._public_key.public_bytes(
-            encoding=serialization.Encoding.Raw,
-            format=serialization.PublicFormat.Raw
-        )
-        return basic_types_pb2.Key(ed25519=public_bytes)
+        return self.to_bytes_raw().hex()
+    
+    def to_string(self) -> str:
+        """
+        Returns the private key as a hex string (raw).
+        """
+        return self.to_string_raw()
 
-    def public_bytes(self, encoding, format):
+
+    def is_ed25519(self) -> bool:
         """
-        Returns the public key bytes in the specified encoding and format.
+        Checks if this public key is Ed25519.
 
-        Args:
-            encoding (Encoding): The encoding to use.
-            format (PublicFormat): The public key format.
+        Returns:
+            bool: True if Ed25519, False otherwise.
+        """
+        return isinstance(self._public_key, ed25519.Ed25519PublicKey)
+
+    def is_ecdsa(self) -> bool:
+        """
+        Checks if this public key is ECDSA (secp256k1).
+
+        Returns:
+            bool: True if ECDSA, False otherwise.
+        """
+        return isinstance(self._public_key, ec.EllipticCurvePublicKey)
+
+    def to_proto(self):
+        """
+        Returns the protobuf representation of the public key.
+        For Ed25519, uses the 'ed25519' field in Key.
+        For ECDSA, uses the 'ECDSASecp256k1' field (may differ by your actual Hedera environment).
 
         Returns:
-            bytes: The public key bytes.
+            Key: The protobuf Key message.
         """
-        return self._public_key.public_bytes(encoding=encoding, format=format)
+        from hedera_sdk_python.hapi.services import basic_types_pb2
+
+        pub_bytes = self.to_bytes_raw()
+        if self.is_ed25519():
+            return basic_types_pb2.Key(ed25519=pub_bytes)
+        else:
+            return basic_types_pb2.Key(ECDSASecp256k1=pub_bytes)
+
+    def __repr__(self):
+        if self.is_ed25519():
+            return f"<PublicKey (Ed25519) hex={self.to_string_raw()}>"
+        return f"<PublicKey (ECDSA) hex={self.to_string_raw()}>"
diff --git a/src/hedera_sdk_python/tokens/token_create_transaction.py b/src/hedera_sdk_python/tokens/token_create_transaction.py
@@ -88,10 +88,7 @@ def build_transaction_body(self):
 
         admin_key_proto = None
         if self.admin_key:
-            admin_public_key_bytes = self.admin_key.public_key().public_bytes(
-                encoding=serialization.Encoding.Raw,
-                format=serialization.PublicFormat.Raw
-            )
+            admin_public_key_bytes = self.admin_key.public_key().to_bytes_raw()
             admin_key_proto = basic_types_pb2.Key(ed25519=admin_public_key_bytes)
 
         token_create_body = token_create_pb2.TokenCreateTransactionBody(
diff --git a/src/hedera_sdk_python/transaction/transaction.py b/src/hedera_sdk_python/transaction/transaction.py
@@ -47,10 +47,7 @@ def sign(self, private_key):
 
         signature = private_key.sign(self.transaction_body_bytes)
 
-        public_key_bytes = private_key.public_key().public_bytes(
-            encoding=Encoding.Raw,
-            format=PublicFormat.Raw
-        )
+        public_key_bytes = private_key.public_key().to_bytes_raw()
 
         sig_pair = basic_types_pb2.SignaturePair(
             pubKeyPrefix=public_key_bytes,
@@ -149,10 +146,7 @@ def is_signed_by(self, public_key):
         Returns:
             bool: True if signed by the given public key, False otherwise.
         """
-        public_key_bytes = public_key.public_bytes(
-            encoding=Encoding.Raw,
-            format=PublicFormat.Raw
-        )
+        public_key_bytes = public_key.to_bytes_raw()
 
         for sig_pair in self.signature_map.sigPair:
             if sig_pair.pubKeyPrefix == public_key_bytes:
diff --git a/tests/test_account_create_transaction.py b/tests/test_account_create_transaction.py
@@ -43,10 +43,7 @@ def test_account_create_transaction_build(mock_account_ids):
 
     transaction_body = account_tx.build_transaction_body()
 
-    expected_public_key_bytes = new_public_key.public_bytes(
-        encoding=serialization.Encoding.Raw,
-        format=serialization.PublicFormat.Raw
-    )
+    expected_public_key_bytes = new_public_key.to_bytes_raw()
 
     assert transaction_body.cryptoCreateAccount.key.ed25519 == expected_public_key_bytes
     assert transaction_body.cryptoCreateAccount.initialBalance == 100000000
diff --git a/tests/test_token_associate_transaction.py b/tests/test_token_associate_transaction.py
@@ -58,7 +58,7 @@ def test_sign_transaction(mock_account_ids):
 
     private_key = MagicMock()
     private_key.sign.return_value = b'signature'
-    private_key.public_key().public_bytes.return_value = b'public_key'
+    private_key.public_key().to_bytes_raw.return_value = b'public_key'
 
     associate_tx.sign(private_key)
 
@@ -80,7 +80,7 @@ def test_to_proto(mock_account_ids):
 
     private_key = MagicMock()
     private_key.sign.return_value = b'signature'
-    private_key.public_key().public_bytes.return_value = b'public_key'
+    private_key.public_key().to_bytes_raw.return_value = b'public_key'
 
     associate_tx.sign(private_key)
     proto = associate_tx.to_proto()
diff --git a/tests/test_token_create_transaction.py b/tests/test_token_create_transaction.py
diff --git a/tests/test_token_delete_transaction.py b/tests/test_token_delete_transaction.py
