Skip to content

Commit d54e9ac

Browse files
philogicaeolethanh
andauthored
Post-SOL fixes (#178)
* Missing chain field on auth * Fix Signature of Solana operation for CRN * Add export_private_key func for accounts * Improve _load_account * Add chain arg to _load_account * Increase default HTTP_REQUEST_TIMEOUT * Typing --------- Co-authored-by: Olivier Le Thanh Duong <[email protected]>
1 parent a636106 commit d54e9ac

File tree

6 files changed

+97
-45
lines changed

6 files changed

+97
-45
lines changed

src/aleph/sdk/account.py

Lines changed: 69 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -10,67 +10,95 @@
1010
from aleph.sdk.chains.remote import RemoteAccount
1111
from aleph.sdk.chains.solana import SOLAccount
1212
from aleph.sdk.conf import load_main_configuration, settings
13+
from aleph.sdk.evm_utils import get_chains_with_super_token
1314
from aleph.sdk.types import AccountFromPrivateKey
1415

1516
logger = logging.getLogger(__name__)
1617

1718
T = TypeVar("T", bound=AccountFromPrivateKey)
1819

20+
chain_account_map: Dict[Chain, Type[T]] = { # type: ignore
21+
Chain.ETH: ETHAccount,
22+
Chain.AVAX: ETHAccount,
23+
Chain.BASE: ETHAccount,
24+
Chain.SOL: SOLAccount,
25+
}
26+
1927

2028
def load_chain_account_type(chain: Chain) -> Type[AccountFromPrivateKey]:
21-
chain_account_map: Dict[Chain, Type[AccountFromPrivateKey]] = {
22-
Chain.ETH: ETHAccount,
23-
Chain.AVAX: ETHAccount,
24-
Chain.SOL: SOLAccount,
25-
Chain.BASE: ETHAccount,
26-
}
27-
return chain_account_map.get(chain) or ETHAccount
29+
return chain_account_map.get(chain) or ETHAccount # type: ignore
2830

2931

30-
def account_from_hex_string(private_key_str: str, account_type: Type[T]) -> T:
32+
def account_from_hex_string(
33+
private_key_str: str,
34+
account_type: Optional[Type[T]],
35+
chain: Optional[Chain] = None,
36+
) -> AccountFromPrivateKey:
3137
if private_key_str.startswith("0x"):
3238
private_key_str = private_key_str[2:]
33-
return account_type(bytes.fromhex(private_key_str))
3439

40+
if not chain:
41+
if not account_type:
42+
account_type = load_chain_account_type(Chain.ETH) # type: ignore
43+
return account_type(bytes.fromhex(private_key_str)) # type: ignore
44+
45+
account_type = load_chain_account_type(chain)
46+
account = account_type(bytes.fromhex(private_key_str))
47+
if chain in get_chains_with_super_token():
48+
account.switch_chain(chain)
49+
return account # type: ignore
3550

36-
def account_from_file(private_key_path: Path, account_type: Type[T]) -> T:
51+
52+
def account_from_file(
53+
private_key_path: Path,
54+
account_type: Optional[Type[T]],
55+
chain: Optional[Chain] = None,
56+
) -> AccountFromPrivateKey:
3757
private_key = private_key_path.read_bytes()
38-
return account_type(private_key)
58+
59+
if not chain:
60+
if not account_type:
61+
account_type = load_chain_account_type(Chain.ETH) # type: ignore
62+
return account_type(private_key) # type: ignore
63+
64+
account_type = load_chain_account_type(chain)
65+
account = account_type(private_key)
66+
if chain in get_chains_with_super_token():
67+
account.switch_chain(chain)
68+
return account
3969

4070

4171
def _load_account(
4272
private_key_str: Optional[str] = None,
4373
private_key_path: Optional[Path] = None,
4474
account_type: Optional[Type[AccountFromPrivateKey]] = None,
75+
chain: Optional[Chain] = None,
4576
) -> AccountFromPrivateKey:
46-
"""Load private key from a string or a file. takes the string argument in priority"""
47-
if private_key_str or (private_key_path and private_key_path.is_file()):
48-
if account_type:
49-
if private_key_path and private_key_path.is_file():
50-
return account_from_file(private_key_path, account_type)
51-
elif private_key_str:
52-
return account_from_hex_string(private_key_str, account_type)
53-
else:
54-
raise ValueError("Any private key specified")
77+
"""Load an account from a private key string or file, or from the configuration file."""
78+
79+
# Loads configuration if no account_type is specified
80+
if not account_type:
81+
config = load_main_configuration(settings.CONFIG_FILE)
82+
if config and hasattr(config, "chain"):
83+
account_type = load_chain_account_type(config.chain)
84+
logger.debug(
85+
f"Detected {config.chain} account for path {settings.CONFIG_FILE}"
86+
)
5587
else:
56-
main_configuration = load_main_configuration(settings.CONFIG_FILE)
57-
if main_configuration:
58-
account_type = load_chain_account_type(main_configuration.chain)
59-
logger.debug(
60-
f"Detected {main_configuration.chain} account for path {settings.CONFIG_FILE}"
61-
)
62-
else:
63-
account_type = ETHAccount # Defaults to ETHAccount
64-
logger.warning(
65-
f"No main configuration data found in {settings.CONFIG_FILE}, defaulting to {account_type.__name__}"
66-
)
67-
if private_key_path and private_key_path.is_file():
68-
return account_from_file(private_key_path, account_type)
69-
elif private_key_str:
70-
return account_from_hex_string(private_key_str, account_type)
71-
else:
72-
raise ValueError("Any private key specified")
88+
account_type = account_type = load_chain_account_type(
89+
Chain.ETH
90+
) # Defaults to ETHAccount
91+
logger.warning(
92+
f"No main configuration data found in {settings.CONFIG_FILE}, defaulting to {account_type and account_type.__name__}"
93+
)
7394

95+
# Loads private key from a string
96+
if private_key_str:
97+
return account_from_hex_string(private_key_str, account_type, chain)
98+
# Loads private key from a file
99+
elif private_key_path and private_key_path.is_file():
100+
return account_from_file(private_key_path, account_type, chain)
101+
# For ledger keys
74102
elif settings.REMOTE_CRYPTO_HOST:
75103
logger.debug("Using remote account")
76104
loop = asyncio.get_event_loop()
@@ -80,10 +108,12 @@ def _load_account(
80108
unix_socket=settings.REMOTE_CRYPTO_UNIX_SOCKET,
81109
)
82110
)
111+
# Fallback: config.path if set, else generate a new private key
83112
else:
84-
account_type = ETHAccount # Defaults to ETHAccount
85113
new_private_key = get_fallback_private_key()
86-
account = account_type(private_key=new_private_key)
114+
account = account_from_hex_string(
115+
bytes.hex(new_private_key), account_type, chain
116+
)
87117
logger.info(
88118
f"Generated fallback private key with address {account.get_address()}"
89119
)

src/aleph/sdk/chains/ethereum.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import asyncio
2+
import base64
23
from decimal import Decimal
34
from pathlib import Path
45
from typing import Awaitable, Optional, Union
@@ -61,6 +62,10 @@ def from_mnemonic(mnemonic: str, chain: Optional[Chain] = None) -> "ETHAccount":
6162
private_key=Account.from_mnemonic(mnemonic=mnemonic).key, chain=chain
6263
)
6364

65+
def export_private_key(self) -> str:
66+
"""Export the private key using standard format."""
67+
return f"0x{base64.b16encode(self.private_key).decode().lower()}"
68+
6469
def get_address(self) -> str:
6570
return self._account.address
6671

src/aleph/sdk/chains/solana.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,12 @@ async def sign_raw(self, buffer: bytes) -> bytes:
4343
sig = self._signing_key.sign(buffer)
4444
return sig.signature
4545

46+
def export_private_key(self) -> str:
47+
"""Export the private key using Phantom format."""
48+
return base58.b58encode(
49+
self.private_key + self._signing_key.verify_key.encode()
50+
).decode()
51+
4652
def get_address(self) -> str:
4753
return encode(self._signing_key.verify_key)
4854

src/aleph/sdk/client/vm_client.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,11 @@
55
from urllib.parse import urlparse
66

77
import aiohttp
8-
from aleph_message.models import ItemHash
8+
from aleph_message.models import Chain, ItemHash
99
from eth_account.messages import encode_defunct
1010
from jwcrypto import jwk
1111

12+
from aleph.sdk.chains.solana import SOLAccount
1213
from aleph.sdk.types import Account
1314
from aleph.sdk.utils import (
1415
create_vm_control_payload,
@@ -36,11 +37,13 @@ def __init__(
3637
self.account = account
3738
self.ephemeral_key = jwk.JWK.generate(kty="EC", crv="P-256")
3839
self.node_url = node_url.rstrip("/")
39-
self.pubkey_payload = self._generate_pubkey_payload()
40+
self.pubkey_payload = self._generate_pubkey_payload(
41+
Chain.SOL if isinstance(account, SOLAccount) else Chain.ETH
42+
)
4043
self.pubkey_signature_header = ""
4144
self.session = session or aiohttp.ClientSession()
4245

43-
def _generate_pubkey_payload(self) -> Dict[str, Any]:
46+
def _generate_pubkey_payload(self, chain: Chain = Chain.ETH) -> Dict[str, Any]:
4447
return {
4548
"pubkey": json.loads(self.ephemeral_key.export_public()),
4649
"alg": "ECDSA",
@@ -50,12 +53,16 @@ def _generate_pubkey_payload(self) -> Dict[str, Any]:
5053
datetime.datetime.utcnow() + datetime.timedelta(days=1)
5154
).isoformat()
5255
+ "Z",
56+
"chain": chain.value,
5357
}
5458

5559
async def _generate_pubkey_signature_header(self) -> str:
5660
pubkey_payload = json.dumps(self.pubkey_payload).encode("utf-8").hex()
57-
signable_message = encode_defunct(hexstr=pubkey_payload)
58-
buffer_to_sign = signable_message.body
61+
if isinstance(self.account, SOLAccount):
62+
buffer_to_sign = bytes(pubkey_payload, encoding="utf-8")
63+
else:
64+
signable_message = encode_defunct(hexstr=pubkey_payload)
65+
buffer_to_sign = signable_message.body
5966

6067
signed_message = await self.account.sign_raw(buffer_to_sign)
6168
pubkey_signature = to_0x_hex(signed_message)

src/aleph/sdk/conf.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ class Settings(BaseSettings):
4141
REMOTE_CRYPTO_HOST: Optional[str] = None
4242
REMOTE_CRYPTO_UNIX_SOCKET: Optional[str] = None
4343
ADDRESS_TO_USE: Optional[str] = None
44-
HTTP_REQUEST_TIMEOUT = 10.0
44+
HTTP_REQUEST_TIMEOUT = 15.0
4545

4646
DEFAULT_CHANNEL: str = "ALEPH-CLOUDSOLUTIONS"
4747
DEFAULT_RUNTIME_ID: str = (

src/aleph/sdk/types.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ def __init__(self, private_key: bytes): ...
3939

4040
async def sign_raw(self, buffer: bytes) -> bytes: ...
4141

42+
def export_private_key(self) -> str: ...
43+
44+
def switch_chain(self, chain: Optional[str] = None) -> None: ...
45+
4246

4347
GenericMessage = TypeVar("GenericMessage", bound=AlephMessage)
4448

0 commit comments

Comments
 (0)