Merge pull request #41 from zksync-sdk/No2FA_pubKeyHash

added ability to supply specific No2FA PubKeyHash(fixes #ZSDK-66)

tests/test_wallet.py

@@ -3,15 +3,15 @@
 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,
                         ZkSyncLibrary, ZkSyncProviderV01, ZkSyncSigner, )
 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"

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):

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

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]

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,47 +32,13 @@ 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',
     # ratio that represents tokens themselves
     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

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)

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

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']