diff --git a/tests/test_wallet.py b/tests/test_wallet.py index 2fd36e7e..2e64be08 100644 --- a/tests/test_wallet.py +++ b/tests/test_wallet.py @@ -3,7 +3,7 @@ from unittest import IsolatedAsyncioTestCase from zksync_sdk.zksync_provider.types import FeeTxType from zksync_sdk.types.responses import Fee - +import asyncio from web3 import Account, HTTPProvider, Web3 from zksync_sdk import (EthereumProvider, EthereumSignerWeb3, HttpJsonRPCTransport, Wallet, ZkSync, @@ -11,7 +11,7 @@ from zksync_sdk.zksync_provider.batch_builder import BatchBuilder from zksync_sdk.network import rinkeby from zksync_sdk.types import ChangePubKeyEcdsa, Token, TransactionWithSignature, \ - TransactionWithOptionalSignature, RatioType, Transfer + TransactionWithOptionalSignature, RatioType, Transfer, AccountTypes from zksync_sdk.zksync_provider.transaction import TransactionStatus from zksync_sdk.wallet import DEFAULT_VALID_FROM, DEFAULT_VALID_UNTIL @@ -348,6 +348,26 @@ async def test_get_tokens(self): async def test_is_signing_key_set(self): assert await self.wallet.is_signing_key_set() + async def test_toggle_2fa(self): + """ + Relate to the server-side code it must be Owned type if enable_2fa is passed + let new_type = if toggle_2fa.enable { + EthAccountType::Owned + } else { + EthAccountType::No2FA + }; + """ + result = await self.wallet.enable_2fa() + self.assertTrue(result) + account_state = await self.wallet.get_account_state() + self.assertEqual(AccountTypes.OWNED, account_state.account_type) + + pub_key_hash = self.wallet.zk_signer.pubkey_hash_str() + result = await self.wallet.disable_2fa(pub_key_hash) + self.assertTrue(result) + account_state = await self.wallet.get_account_state() + self.assertEqual(AccountTypes.NO_2FA, account_state.account_type) + class TestEthereumProvider(IsolatedAsyncioTestCase): private_key = "0xa045b52470d306ff78e91b0d2d92f90f7504189125a46b69423dc673fd6b4f3e" diff --git a/zksync_sdk/types/__init__.py b/zksync_sdk/types/__init__.py index a93784c0..2f421ba2 100644 --- a/zksync_sdk/types/__init__.py +++ b/zksync_sdk/types/__init__.py @@ -3,6 +3,7 @@ from .responses import * from .signatures import * from .transactions import * +from .auth_types import * class ChainId(IntEnum): diff --git a/zksync_sdk/types/auth_types.py b/zksync_sdk/types/auth_types.py new file mode 100644 index 00000000..523d18c3 --- /dev/null +++ b/zksync_sdk/types/auth_types.py @@ -0,0 +1,84 @@ +from enum import Enum +from dataclasses import dataclass +from typing import Optional +from zksync_sdk.types.signatures import TxEthSignature + + +class ChangePubKeyTypes(Enum): + onchain = "Onchain" + ecdsa = "ECDSA" + create2 = "CREATE2" + + +@dataclass +class ChangePubKeyEcdsa: + batch_hash: bytes = b"\x00" * 32 + + def encode_message(self) -> bytes: + return self.batch_hash + + def dict(self, signature: str): + return {"type": "ECDSA", + "ethSignature": signature, + "batchHash": f"0x{self.batch_hash.hex()}"} + + +@dataclass +class ChangePubKeyCREATE2: + creator_address: str + salt_arg: bytes + code_hash: bytes + + def encode_message(self) -> bytes: + return self.salt_arg + + def dict(self): + return {"type": "CREATE2", + "saltArg": f"0x{self.salt_arg.hex()}", + "codeHash": f"0x{self.code_hash.hex()}"} + + +@dataclass +class Toggle2FA: + enable: bool + account_id: int + time_stamp_milliseconds: int + signature: TxEthSignature + pub_key_hash: Optional[str] + + def dict(self): + if self.pub_key_hash is not None: + return { + "enable": self.enable, + "accountId": self.account_id, + "timestamp": self.time_stamp_milliseconds, + "signature": self.signature.dict(), + "pubKeyHash": self.pub_key_hash + } + else: + return { + "enable": self.enable, + "accountId": self.account_id, + "timestamp": self.time_stamp_milliseconds, + "signature": self.signature.dict(), + } + + +def get_toggle_message(require_2fa: bool, time_stamp: int) -> str: + if require_2fa: + msg = f"By signing this message, you are opting into Two-factor Authentication protection by the zkSync " \ + f"Server.\n" \ + f"Transactions now require signatures by both your L1 and L2 private key.\n" \ + f"Timestamp: {time_stamp}" + else: + msg = f"You are opting out of Two-factor Authentication protection by the zkSync Server.\n" \ + f"Transactions now only require signatures by your L2 private key.\n" \ + f"BY SIGNING THIS MESSAGE, YOU ARE TRUSTING YOUR WALLET CLIENT TO KEEP YOUR L2 PRIVATE KEY SAFE!\n" \ + f"Timestamp: {time_stamp}" + return msg + + +def get_toggle_message_with_pub(require_2fa: bool, time_stamp: int, pub_key_hash: str) -> str: + msg = get_toggle_message(require_2fa, time_stamp) + msg += f"\nPubKeyHash: {pub_key_hash}" + return msg diff --git a/zksync_sdk/types/responses.py b/zksync_sdk/types/responses.py index 31c61bf0..7782d6da 100644 --- a/zksync_sdk/types/responses.py +++ b/zksync_sdk/types/responses.py @@ -1,4 +1,5 @@ from typing import Any, Dict, Optional +from enum import Enum from decimal import Decimal from zksync_sdk.types.transactions import Token @@ -47,9 +48,16 @@ class Config: alias_generator = to_camel +class AccountTypes(str, Enum): + OWNED = "Owned", + CREATE2 = "CREATE2", + NO_2FA = "No2FA" + + class AccountState(BaseModel): address: str id: Optional[int] + account_type: Optional[AccountTypes] depositing: Optional[Depositing] committed: Optional[State] verified: Optional[State] diff --git a/zksync_sdk/types/transactions.py b/zksync_sdk/types/transactions.py index ef51dc3d..ba58c996 100644 --- a/zksync_sdk/types/transactions.py +++ b/zksync_sdk/types/transactions.py @@ -13,6 +13,7 @@ serialize_nonce, serialize_timestamp, serialize_token_id, serialize_ratio_part) from zksync_sdk.types.signatures import TxEthSignature, TxSignature +from zksync_sdk.types.auth_types import ChangePubKeyCREATE2, ChangePubKeyEcdsa DEFAULT_TOKEN_ADDRESS = "0x0000000000000000000000000000000000000000" @@ -31,12 +32,6 @@ class EncodedTxType(IntEnum): WITHDRAW_NFT = 10 -class ChangePubKeyTypes(Enum): - onchain = "Onchain" - ecdsa = "ECDSA" - create2 = "CREATE2" - - class RatioType(Enum): # ratio that represents the lowest denominations of tokens (wei for ETH, satoshi for BTC etc.) wei = 'Wei', @@ -44,34 +39,6 @@ class RatioType(Enum): token = 'Token' -@dataclass -class ChangePubKeyEcdsa: - batch_hash: bytes = b"\x00" * 32 - - def encode_message(self): - return self.batch_hash - - def dict(self, signature: str): - return {"type": "ECDSA", - "ethSignature": signature, - "batchHash": f"0x{self.batch_hash.hex()}"} - - -@dataclass -class ChangePubKeyCREATE2: - creator_address: str - salt_arg: bytes - code_hash: bytes - - def encode_message(self): - return self.salt_arg - - def dict(self): - return {"type": "CREATE2", - "saltArg": f"0x{self.salt_arg.hex()}", - "codeHash": f"0x{self.code_hash.hex()}"} - - class Token(BaseModel): address: str id: int @@ -111,11 +78,13 @@ def decimal_str_amount(self, amount: int) -> str: return d_str + def token_ratio_to_wei_ratio(token_ratio: Fraction, token_sell: Token, token_buy: Token) -> Fraction: num = token_sell.from_decimal(Decimal(token_ratio.numerator)) den = token_buy.from_decimal(Decimal(token_ratio.denominator)) return Fraction(num, den) + class Tokens(BaseModel): tokens: List[Token] @@ -561,7 +530,6 @@ def get_sig(opt_value: Optional[TxEthSignature]) -> TxEthSignature: return w3.eth.account.recover_message(encoded_message, signature=sig) - @dataclass class Swap(EncodedTx): submitter_id: int diff --git a/zksync_sdk/wallet.py b/zksync_sdk/wallet.py index 42be5b09..7cea6d9a 100644 --- a/zksync_sdk/wallet.py +++ b/zksync_sdk/wallet.py @@ -1,3 +1,4 @@ +import time from decimal import Decimal from fractions import Fraction from typing import List, Optional, Tuple, Union @@ -7,7 +8,8 @@ from zksync_sdk.types import (ChangePubKey, ChangePubKeyCREATE2, ChangePubKeyEcdsa, ChangePubKeyTypes, EncodedTx, ForcedExit, Token, TokenLike, Tokens, TransactionWithSignature, Transfer, TxEthSignature, - Withdraw, MintNFT, WithdrawNFT, NFT, Order, Swap, RatioType, token_ratio_to_wei_ratio) + Withdraw, MintNFT, WithdrawNFT, NFT, Order, Swap, RatioType, + token_ratio_to_wei_ratio, get_toggle_message, get_toggle_message_with_pub, Toggle2FA) from zksync_sdk.zksync_provider import FeeTxType, ZkSyncProviderInterface from zksync_sdk.zksync_signer import ZkSyncSigner from zksync_sdk.zksync_provider.transaction import Transaction @@ -491,3 +493,35 @@ async def resolve_token(self, token: TokenLike) -> Token: if resolved_token is None: raise TokenNotFoundError return resolved_token + + async def enable_2fa(self) -> bool: + mil_seconds = int(time.time() * 1000) + msg = get_toggle_message(True, mil_seconds) + eth_sig = self.eth_signer.sign(msg.encode()) + account_id = await self.get_account_id() + toggle = Toggle2FA(True, + account_id, + mil_seconds, + eth_sig, + None + ) + return await self.zk_provider.toggle_2fa(toggle) + + async def disable_2fa(self, pub_key_hash: Optional[str]) -> bool: + mil_seconds = int(time.time() * 1000) + if pub_key_hash is None: + msg = get_toggle_message(False, mil_seconds) + else: + msg = get_toggle_message_with_pub(False, mil_seconds, pub_key_hash) + eth_sig = self.eth_signer.sign(msg.encode()) + account_id = await self.get_account_id() + toggle = Toggle2FA(False, + account_id, + mil_seconds, + eth_sig, + pub_key_hash) + return await self.zk_provider.toggle_2fa(toggle) + + async def disable_2fa_with_pub_key(self): + pub_key_hash = self.zk_signer.pubkey_hash_str() + return await self.disable_2fa(pub_key_hash) diff --git a/zksync_sdk/zksync_provider/interface.py b/zksync_sdk/zksync_provider/interface.py index 94c2b946..fd5292b4 100644 --- a/zksync_sdk/zksync_provider/interface.py +++ b/zksync_sdk/zksync_provider/interface.py @@ -1,14 +1,11 @@ from abc import ABC, abstractmethod from decimal import Decimal -from typing import List, Optional, Tuple, Union - -from eth_typing import Address - +from typing import List, Optional, Union from zksync_sdk.transport import JsonRPCTransport from zksync_sdk.types import (AccountState, ContractAddress, EncodedTx, EthOpInfo, Fee, Token, TokenLike, Tokens, TransactionDetails, TransactionWithSignature, TransactionWithOptionalSignature, - TxEthSignature, ) + TxEthSignature, Toggle2FA, ) from zksync_sdk.zksync_provider.types import FeeTxType from zksync_sdk.zksync_provider.transaction import Transaction @@ -81,3 +78,7 @@ async def get_transaction_fee(self, tx_type: FeeTxType, address: str, @abstractmethod async def get_token_price(self, token: Token) -> Decimal: raise NotImplementedError + + @abstractmethod + async def toggle_2fa(self, toggle2fa: Toggle2FA) -> bool: + raise NotImplementedError diff --git a/zksync_sdk/zksync_provider/v01.py b/zksync_sdk/zksync_provider/v01.py index 7939083d..a539bcbb 100644 --- a/zksync_sdk/zksync_provider/v01.py +++ b/zksync_sdk/zksync_provider/v01.py @@ -1,13 +1,12 @@ +from dataclasses import asdict from decimal import Decimal -from typing import List, Optional, Tuple, Union - -from eth_typing import Address +from typing import List, Optional, Union from web3 import Web3 from zksync_sdk.types import (AccountState, ContractAddress, EncodedTx, EthOpInfo, Fee, Token, TokenLike, Tokens, TransactionDetails, TransactionWithSignature, TransactionWithOptionalSignature, - TxEthSignature, ) + TxEthSignature, Toggle2FA, ) from zksync_sdk.zksync_provider.error import AccountDoesNotExist from zksync_sdk.zksync_provider.interface import ZkSyncProviderInterface from zksync_sdk.zksync_provider.types import FeeTxType @@ -64,6 +63,9 @@ async def get_state(self, address: str) -> AccountState: data = await self.provider.request("account_info", [address]) if data is None: raise AccountDoesNotExist(address=address) + if "accountType" in data and isinstance(data["accountType"], dict) and \ + list(data["accountType"].keys())[0] == 'No2FA': + data["accountType"] = 'No2FA' return AccountState(**data) async def get_confirmations_for_eth_op_amount(self) -> int: @@ -102,3 +104,7 @@ async def get_transaction_fee(self, tx_type: FeeTxType, address: str, async def get_token_price(self, token: Token) -> Decimal: data = await self.provider.request('get_token_price', [token.symbol]) return Decimal(data) + + async def toggle_2fa(self, toggle2fa: Toggle2FA) -> bool: + data = await self.provider.request('toggle_2fa', [toggle2fa.dict()]) + return 'success' in data and data['success']