Skip to content

Commit 2790df9

Browse files
1yamnesitor
andauthored
Feature: Account Handler (#175)
* Feature: Internal account management + fix on _load_account to handle SolAccount * fixup! Feature: Internal account management + fix on _load_account to handle SolAccount * Fix: chains_config wasn't using settings.CONFIG_HOME for locations * Fix: blakc issue * Fix: rename CHAINS_CONFIG_FILE to CONFIG_FILE to avoid getting issue by conf of chain * Fix: base58 and pynacl is now needed for build * Fix: f string without nay placeholders * Fix: black error * Refactor: we now store single account at the time * Fix: ruff issue * fix: debug stuff remove * Fix: Improve code structure in pair-programming with Lyam --------- Co-authored-by: Andres D. Molins <[email protected]>
1 parent cf70462 commit 2790df9

File tree

6 files changed

+263
-15
lines changed

6 files changed

+263
-15
lines changed

.gitignore

+1-1
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ MANIFEST
5050
**/device.key
5151

5252
# environment variables
53-
.env
53+
.config.json
5454
.env.local
5555

5656
.gitsigners

pyproject.toml

+2
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ dependencies = [
3434
"aleph-superfluid>=0.2.1",
3535
"eth_typing==4.3.1",
3636
"web3==6.3.0",
37+
"base58==2.1.1", # Needed now as default with _load_account changement
38+
"pynacl==1.5.0" # Needed now as default with _load_account changement
3739
]
3840

3941
[project.optional-dependencies]

src/aleph/sdk/account.py

+43-9
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,32 @@
11
import asyncio
22
import logging
33
from pathlib import Path
4-
from typing import Optional, Type, TypeVar
4+
from typing import Dict, Optional, Type, TypeVar
5+
6+
from aleph_message.models import Chain
57

68
from aleph.sdk.chains.common import get_fallback_private_key
79
from aleph.sdk.chains.ethereum import ETHAccount
810
from aleph.sdk.chains.remote import RemoteAccount
9-
from aleph.sdk.conf import settings
11+
from aleph.sdk.chains.solana import SOLAccount
12+
from aleph.sdk.conf import load_main_configuration, settings
1013
from aleph.sdk.types import AccountFromPrivateKey
1114

1215
logger = logging.getLogger(__name__)
1316

1417
T = TypeVar("T", bound=AccountFromPrivateKey)
1518

1619

20+
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
28+
29+
1730
def account_from_hex_string(private_key_str: str, account_type: Type[T]) -> T:
1831
if private_key_str.startswith("0x"):
1932
private_key_str = private_key_str[2:]
@@ -28,16 +41,36 @@ def account_from_file(private_key_path: Path, account_type: Type[T]) -> T:
2841
def _load_account(
2942
private_key_str: Optional[str] = None,
3043
private_key_path: Optional[Path] = None,
31-
account_type: Type[AccountFromPrivateKey] = ETHAccount,
44+
account_type: Optional[Type[AccountFromPrivateKey]] = None,
3245
) -> AccountFromPrivateKey:
3346
"""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")
55+
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")
3473

35-
if private_key_str:
36-
logger.debug("Using account from string")
37-
return account_from_hex_string(private_key_str, account_type)
38-
elif private_key_path and private_key_path.is_file():
39-
logger.debug("Using account from file")
40-
return account_from_file(private_key_path, account_type)
4174
elif settings.REMOTE_CRYPTO_HOST:
4275
logger.debug("Using remote account")
4376
loop = asyncio.get_event_loop()
@@ -48,6 +81,7 @@ def _load_account(
4881
)
4982
)
5083
else:
84+
account_type = ETHAccount # Defaults to ETHAccount
5185
new_private_key = get_fallback_private_key()
5286
account = account_type(private_key=new_private_key)
5387
logger.info(

src/aleph/sdk/chains/solana.py

+91-3
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import json
22
from pathlib import Path
3-
from typing import Dict, Optional, Union
3+
from typing import Dict, List, Optional, Union
44

55
import base58
66
from nacl.exceptions import BadSignatureError as NaclBadSignatureError
@@ -22,7 +22,7 @@ class SOLAccount(BaseAccount):
2222
_private_key: PrivateKey
2323

2424
def __init__(self, private_key: bytes):
25-
self.private_key = private_key
25+
self.private_key = parse_private_key(private_key_from_bytes(private_key))
2626
self._signing_key = SigningKey(self.private_key)
2727
self._private_key = self._signing_key.to_curve25519_private_key()
2828

@@ -79,7 +79,7 @@ def verify_signature(
7979
public_key: The public key to use for verification. Can be a base58 encoded string or bytes.
8080
message: The message to verify. Can be an utf-8 string or bytes.
8181
Raises:
82-
BadSignatureError: If the signature is invalid.
82+
BadSignatureError: If the signature is invalid.!
8383
"""
8484
if isinstance(signature, str):
8585
signature = base58.b58decode(signature)
@@ -91,3 +91,91 @@ def verify_signature(
9191
VerifyKey(public_key).verify(message, signature)
9292
except NaclBadSignatureError as e:
9393
raise BadSignatureError from e
94+
95+
96+
def private_key_from_bytes(
97+
private_key_bytes: bytes, output_format: str = "base58"
98+
) -> Union[str, List[int], bytes]:
99+
"""
100+
Convert a Solana private key in bytes back to different formats (base58 string, uint8 list, or raw bytes).
101+
102+
- For base58 string: Encode the bytes into a base58 string.
103+
- For uint8 list: Convert the bytes into a list of integers.
104+
- For raw bytes: Return as-is.
105+
106+
Args:
107+
private_key_bytes (bytes): The private key in byte format.
108+
output_format (str): The format to return ('base58', 'list', 'bytes').
109+
110+
Returns:
111+
The private key in the requested format.
112+
113+
Raises:
114+
ValueError: If the output_format is not recognized or the private key length is invalid.
115+
"""
116+
if not isinstance(private_key_bytes, bytes):
117+
raise ValueError("Expected the private key in bytes.")
118+
119+
if len(private_key_bytes) != 32:
120+
raise ValueError("Solana private key must be exactly 32 bytes long.")
121+
122+
if output_format == "base58":
123+
return base58.b58encode(private_key_bytes).decode("utf-8")
124+
125+
elif output_format == "list":
126+
return list(private_key_bytes)
127+
128+
elif output_format == "bytes":
129+
return private_key_bytes
130+
131+
else:
132+
raise ValueError("Invalid output format. Choose 'base58', 'list', or 'bytes'.")
133+
134+
135+
def parse_private_key(private_key: Union[str, List[int], bytes]) -> bytes:
136+
"""
137+
Parse the private key which could be either:
138+
- a base58-encoded string (which may contain both private and public key)
139+
- a list of uint8 integers (which may contain both private and public key)
140+
- a byte array (exactly 32 bytes)
141+
142+
Returns:
143+
bytes: The private key in byte format (32 bytes).
144+
145+
Raises:
146+
ValueError: If the private key format is invalid or the length is incorrect.
147+
"""
148+
# If the private key is already in byte format
149+
if isinstance(private_key, bytes):
150+
if len(private_key) != 32:
151+
raise ValueError("The private key in bytes must be exactly 32 bytes long.")
152+
return private_key
153+
154+
# If the private key is a base58-encoded string
155+
elif isinstance(private_key, str):
156+
try:
157+
decoded_key = base58.b58decode(private_key)
158+
if len(decoded_key) not in [32, 64]:
159+
raise ValueError(
160+
"The base58 decoded private key must be either 32 or 64 bytes long."
161+
)
162+
return decoded_key[:32]
163+
except Exception as e:
164+
raise ValueError(f"Invalid base58 encoded private key: {e}")
165+
166+
# If the private key is a list of uint8 integers
167+
elif isinstance(private_key, list):
168+
if all(isinstance(i, int) and 0 <= i <= 255 for i in private_key):
169+
byte_key = bytes(private_key)
170+
if len(byte_key) < 32:
171+
raise ValueError("The uint8 array must contain at least 32 elements.")
172+
return byte_key[:32] # Take the first 32 bytes (private key)
173+
else:
174+
raise ValueError(
175+
"Invalid uint8 array, must contain integers between 0 and 255."
176+
)
177+
178+
else:
179+
raise ValueError(
180+
"Unsupported private key format. Must be a base58 string, bytes, or a list of uint8 integers."
181+
)

src/aleph/sdk/conf.py

+67-1
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,27 @@
1+
import json
2+
import logging
13
import os
24
from pathlib import Path
35
from shutil import which
46
from typing import Dict, Optional, Union
57

68
from aleph_message.models import Chain
79
from aleph_message.models.execution.environment import HypervisorType
8-
from pydantic import BaseSettings, Field
10+
from pydantic import BaseModel, BaseSettings, Field
911

1012
from aleph.sdk.types import ChainInfo
1113

14+
logger = logging.getLogger(__name__)
15+
1216

1317
class Settings(BaseSettings):
1418
CONFIG_HOME: Optional[str] = None
1519

20+
CONFIG_FILE: Path = Field(
21+
default=Path("config.json"),
22+
description="Path to the JSON file containing chain account configurations",
23+
)
24+
1625
# In case the user does not want to bother with handling private keys himself,
1726
# do an ugly and insecure write and read from disk to this file.
1827
PRIVATE_KEY_FILE: Path = Field(
@@ -139,6 +148,18 @@ class Config:
139148
env_file = ".env"
140149

141150

151+
class MainConfiguration(BaseModel):
152+
"""
153+
Intern Chain Management with Account.
154+
"""
155+
156+
path: Path
157+
chain: Chain
158+
159+
class Config:
160+
use_enum_values = True
161+
162+
142163
# Settings singleton
143164
settings = Settings()
144165

@@ -162,6 +183,19 @@ class Config:
162183
settings.PRIVATE_MNEMONIC_FILE = Path(
163184
settings.CONFIG_HOME, "private-keys", "substrate.mnemonic"
164185
)
186+
if str(settings.CONFIG_FILE) == "config.json":
187+
settings.CONFIG_FILE = Path(settings.CONFIG_HOME, "config.json")
188+
# If Config file exist and well filled we update the PRIVATE_KEY_FILE default
189+
if settings.CONFIG_FILE.exists():
190+
try:
191+
with open(settings.CONFIG_FILE, "r", encoding="utf-8") as f:
192+
config_data = json.load(f)
193+
194+
if "path" in config_data:
195+
settings.PRIVATE_KEY_FILE = Path(config_data["path"])
196+
except json.JSONDecodeError:
197+
pass
198+
165199

166200
# Update CHAINS settings and remove placeholders
167201
CHAINS_ENV = [(key[7:], value) for key, value in settings if key.startswith("CHAINS_")]
@@ -172,3 +206,35 @@ class Config:
172206
field = field.lower()
173207
settings.CHAINS[chain].__dict__[field] = value
174208
settings.__delattr__(f"CHAINS_{fields}")
209+
210+
211+
def save_main_configuration(file_path: Path, data: MainConfiguration):
212+
"""
213+
Synchronously save a single ChainAccount object as JSON to a file.
214+
"""
215+
with file_path.open("w") as file:
216+
data_serializable = data.dict()
217+
data_serializable["path"] = str(data_serializable["path"])
218+
json.dump(data_serializable, file, indent=4)
219+
220+
221+
def load_main_configuration(file_path: Path) -> Optional[MainConfiguration]:
222+
"""
223+
Synchronously load the private key and chain type from a file.
224+
If the file does not exist or is empty, return None.
225+
"""
226+
if not file_path.exists() or file_path.stat().st_size == 0:
227+
logger.debug(f"File {file_path} does not exist or is empty. Returning None.")
228+
return None
229+
230+
try:
231+
with file_path.open("rb") as file:
232+
content = file.read()
233+
data = json.loads(content.decode("utf-8"))
234+
return MainConfiguration(**data)
235+
except UnicodeDecodeError as e:
236+
logger.error(f"Unable to decode {file_path} as UTF-8: {e}")
237+
except json.JSONDecodeError:
238+
logger.error(f"Invalid JSON format in {file_path}.")
239+
240+
return None

tests/unit/test_chain_solana.py

+59-1
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,12 @@
88
from nacl.signing import VerifyKey
99

1010
from aleph.sdk.chains.common import get_verification_buffer
11-
from aleph.sdk.chains.solana import SOLAccount, get_fallback_account, verify_signature
11+
from aleph.sdk.chains.solana import (
12+
SOLAccount,
13+
get_fallback_account,
14+
parse_private_key,
15+
verify_signature,
16+
)
1217
from aleph.sdk.exceptions import BadSignatureError
1318

1419

@@ -136,3 +141,56 @@ async def test_sign_raw(solana_account):
136141
assert isinstance(signature, bytes)
137142

138143
verify_signature(signature, solana_account.get_address(), buffer)
144+
145+
146+
def test_parse_solana_private_key_bytes():
147+
# Valid 32-byte private key
148+
private_key_bytes = bytes(range(32))
149+
parsed_key = parse_private_key(private_key_bytes)
150+
assert isinstance(parsed_key, bytes)
151+
assert len(parsed_key) == 32
152+
assert parsed_key == private_key_bytes
153+
154+
# Invalid private key (too short)
155+
with pytest.raises(
156+
ValueError, match="The private key in bytes must be exactly 32 bytes long."
157+
):
158+
parse_private_key(bytes(range(31)))
159+
160+
161+
def test_parse_solana_private_key_base58():
162+
# Valid base58 private key (32 bytes)
163+
base58_key = base58.b58encode(bytes(range(32))).decode("utf-8")
164+
parsed_key = parse_private_key(base58_key)
165+
assert isinstance(parsed_key, bytes)
166+
assert len(parsed_key) == 32
167+
168+
# Invalid base58 key (not decodable)
169+
with pytest.raises(ValueError, match="Invalid base58 encoded private key"):
170+
parse_private_key("invalid_base58_key")
171+
172+
# Invalid base58 key (wrong length)
173+
with pytest.raises(
174+
ValueError,
175+
match="The base58 decoded private key must be either 32 or 64 bytes long.",
176+
):
177+
parse_private_key(base58.b58encode(bytes(range(31))).decode("utf-8"))
178+
179+
180+
def test_parse_solana_private_key_list():
181+
# Valid list of uint8 integers (64 elements, but we only take the first 32 for private key)
182+
uint8_list = list(range(64))
183+
parsed_key = parse_private_key(uint8_list)
184+
assert isinstance(parsed_key, bytes)
185+
assert len(parsed_key) == 32
186+
assert parsed_key == bytes(range(32))
187+
188+
# Invalid list (contains non-integers)
189+
with pytest.raises(ValueError, match="Invalid uint8 array"):
190+
parse_private_key([1, 2, "not an int", 4]) # type: ignore # Ignore type check for string
191+
192+
# Invalid list (less than 32 elements)
193+
with pytest.raises(
194+
ValueError, match="The uint8 array must contain at least 32 elements."
195+
):
196+
parse_private_key(list(range(31)))

0 commit comments

Comments
 (0)