thirdweb-dev · adam-maj · Jul 27, 2022 · Jul 27, 2022 · Jul 27, 2022
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -6,7 +6,7 @@ packages = [
   {include = "thirdweb"},
 ]
 readme = "README.md"
-version = "2.0.17"
+version = "2.0.18"
 
 [tool.poetry.dependencies]
 python = ">=3.7.1"
@@ -18,6 +18,7 @@ thirdweb-eth-account = "^0.6.6"
 pymerkle = "^2.0.2"
 pyee = "^9.0.4"
 cbor2 = "^5.4.3"
+pytz = "^2022.1"
 
 [tool.poetry.dev-dependencies]
 black = "^22.1.0"

diff --git a/tests/test_auth.py b/tests/test_auth.py
@@ -23,7 +23,7 @@ def test_verify_with_params(sdk: ThirdwebSDK, primary_account, secondary_account
     payload = sdk.auth.login(
         domain,
         LoginOptions(
-            expirationTime=datetime.utcnow() + timedelta(hours=5), chainId=137
+            expiration_time=datetime.utcnow() + timedelta(hours=5), chain_id=137
         ),
     )
 
@@ -32,7 +32,7 @@ def test_verify_with_params(sdk: ThirdwebSDK, primary_account, secondary_account
         domain,
         payload,
         VerifyOptions(
-            chainId=137,
+            chain_id=137,
         ),
     )
 
@@ -56,7 +56,7 @@ def test_reject_incorrect_domain(sdk: ThirdwebSDK, primary_account, secondary_ac
 def test_reject_expired_payload(sdk: ThirdwebSDK, primary_account, secondary_account):
     sdk.update_signer(primary_account)
     payload = sdk.auth.login(
-        domain, LoginOptions(expirationTime=datetime.utcnow() - timedelta(hours=5))
+        domain, LoginOptions(expiration_time=datetime.utcnow() - timedelta(hours=5))
     )
 
     sdk.update_signer(secondary_account)
@@ -72,11 +72,11 @@ def test_reject_incorrect_chain_id(
     sdk: ThirdwebSDK, primary_account, secondary_account
 ):
     sdk.update_signer(primary_account)
-    payload = sdk.auth.login(domain, LoginOptions(chainId=1))
+    payload = sdk.auth.login(domain, LoginOptions(chain_id=1))
 
     sdk.update_signer(secondary_account)
     try:
-        sdk.auth.verify(domain, payload, VerifyOptions(chainId=137))
+        sdk.auth.verify(domain, payload, VerifyOptions(chain_id=137))
         assert False
     except Exception as e:
         assert "Chain ID '137' does not match payload chain ID '1'" in str(e)
@@ -138,7 +138,7 @@ def test_reject_token_before_invalid(
     token = sdk.auth.generate_auth_token(
         domain,
         payload,
-        AuthenticationOptions(invalidBefore=datetime.utcnow() + timedelta(hours=5)),
+        AuthenticationOptions(invalid_before=datetime.utcnow() + timedelta(hours=5)),
     )
     try:
         sdk.auth.authenticate(domain, token)
@@ -156,7 +156,7 @@ def test_reject_expired_token(sdk: ThirdwebSDK, primary_account, secondary_accou
     token = sdk.auth.generate_auth_token(
         domain,
         payload,
-        AuthenticationOptions(expirationTime=datetime.utcnow() - timedelta(hours=5)),
+        AuthenticationOptions(expiration_time=datetime.utcnow() - timedelta(hours=5)),
     )
     try:
         sdk.auth.authenticate(domain, token)

diff --git a/thirdweb/core/auth/wallet_authenticator.py b/thirdweb/core/auth/wallet_authenticator.py
@@ -17,6 +17,7 @@
 from thirdweb.types.sdk import SDKOptions
 from eth_account.messages import encode_defunct
 from eth_account.datastructures import SignedMessage
+import pytz
 
 
 class WalletAuthenticator(ProviderHandler):
@@ -80,12 +81,12 @@ def login(
         signer_address = self._require_signer().address
         payload_data = LoginPayloadData(
             domain=domain,
-            expirationTime=options.expirationTime
-            if options.expirationTime is not None
+            expiration_time=options.expiration_time
+            if options.expiration_time is not None
             else datetime.utcnow() + timedelta(minutes=5),
             address=signer_address,
             nonce=options.nonce if options.nonce is not None else str(uuid4()),
-            chainId=options.chainId,
+            chain_id=options.chain_id,
         )
 
         message = self._generate_message(payload_data)
@@ -127,13 +128,18 @@ def verify(
 
         # Check that the payload hasn't expired
         current_time = datetime.utcnow()
-        if current_time > payload.payload.expirationTime:
+        if current_time.replace(
+            tzinfo=pytz.utc
+        ) > payload.payload.expiration_time.replace(tzinfo=pytz.utc):
             raise Exception(f"Login request has expired")
 
         # If chain ID is specified, check that it matches the chain ID of the signature
-        if options.chainId is not None and options.chainId != payload.payload.chainId:
+        if (
+            options.chain_id is not None
+            and options.chain_id != payload.payload.chain_id
+        ):
             raise Exception(
-                f"Chain ID '{options.chainId}' does not match payload chain ID '{payload.payload.chainId}'"
+                f"Chain ID '{options.chain_id}' does not match payload chain ID '{payload.payload.chain_id}'"
             )
 
         # Check that the signing address is the claimed wallet address
@@ -176,11 +182,11 @@ def generate_auth_token(
             iss=admin_address,
             sub=user_address,
             aud=domain,
-            nbf=int(options.invalidBefore.timestamp())
-            if options.invalidBefore is not None
+            nbf=int(options.invalid_before.timestamp())
+            if options.invalid_before is not None
             else int(datetime.utcnow().timestamp()),
-            exp=int(options.expirationTime.timestamp())
-            if options.expirationTime is not None
+            exp=int(options.expiration_time.timestamp())
+            if options.expiration_time is not None
             else int((datetime.utcnow() + timedelta(hours=5)).timestamp()),
             iat=int(datetime.utcnow().timestamp()),
             jti=str(uuid4()),
@@ -243,13 +249,17 @@ def authenticate(
 
         # Check that the token is past the invalid before time
         current_time = datetime.utcnow()
-        if current_time < datetime.fromtimestamp(payload.nbf):
+        if current_time.replace(tzinfo=pytz.utc) < datetime.fromtimestamp(
+            payload.nbf
+        ).replace(tzinfo=pytz.utc):
             raise Exception(
                 f"This token is invalid before epoch time '{payload.nbf}', current epoch time is '{int(current_time.timestamp())}'"
             )
 
         # Check that the token hasn't expired
-        if current_time > datetime.fromtimestamp(payload.exp):
+        if current_time.replace(tzinfo=pytz.utc) > datetime.fromtimestamp(
+            payload.exp
+        ).replace(tzinfo=pytz.utc):
             raise Exception(
                 f"This token expired at epoch time '{payload.exp}', current epoch time is '{int(current_time.timestamp())}'"
             )
@@ -289,11 +299,15 @@ def _generate_message(self, payload: LoginPayloadData) -> str:
         message += "Make sure that the requesting domain above matches the URL of the current website.\n\n"
 
         # Add data fields in compliance with the EIP-4361 standard
-        if payload.chainId is not None:
-            message += f"Chain ID: {payload.chainId}\n"
+        if payload.chain_id is not None:
+            message += f"Chain ID: {payload.chain_id}\n"
 
         message += f"Nonce: {payload.nonce}\n"
-        message += f"Expiration Time: {payload.expirationTime.isoformat()}\n"
+
+        time = payload.expiration_time.strftime("%Y-%m-%dT%H:%M:%S")
+        microseconds = payload.expiration_time.strftime("%f")[0:3]
+        formatted_time = f"{time}.{microseconds}Z"
+        message += f"Expiration Time: {formatted_time}\n"
 
         return message
 

diff --git a/thirdweb/types/__init__.py b/thirdweb/types/__init__.py
@@ -3,3 +3,4 @@
 from .multiwrap import *
 from .marketplace import *
 from .settings import *
+from .auth import *
diff --git a/thirdweb/types/auth.py b/thirdweb/types/auth.py
@@ -1,39 +1,57 @@
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional
+from typing import Any, Dict, Optional
 
 
 @dataclass
 class LoginOptions:
     nonce: Optional[str] = None
-    expirationTime: Optional[datetime] = None
-    chainId: Optional[int] = None
+    expiration_time: Optional[datetime] = None
+    chain_id: Optional[int] = None
 
 
 @dataclass
 class LoginPayloadData:
     domain: str
     address: str
     nonce: str
-    expirationTime: datetime
-    chainId: Optional[int] = None
+    expiration_time: datetime
+    chain_id: Optional[int] = None
+
+    @staticmethod
+    def from_json(json: Dict[str, Any]) -> "LoginPayloadData":
+        return LoginPayloadData(
+            json["domain"],
+            json["address"],
+            json["nonce"],
+            datetime.strptime(json["expiration_time"], "%Y-%m-%dT%H:%M:%S.%fZ"),
+            json.get("chain_id"),
+        )
 
 
 @dataclass
 class LoginPayload:
     payload: LoginPayloadData
     signature: str
 
+    @staticmethod
+    def from_json(json: Dict[str, Any]) -> "LoginPayload":
+        data = LoginPayloadData.from_json(json["payload"])
+        return LoginPayload(
+            data,
+            signature=json["signature"],
+        )
+
 
 @dataclass
 class VerifyOptions:
-    chainId: Optional[int] = None
+    chain_id: Optional[int] = None
 
 
 @dataclass
 class AuthenticationOptions:
-    invalidBefore: Optional[datetime] = None
-    expirationTime: Optional[datetime] = None
+    invalid_before: Optional[datetime] = None
+    expiration_time: Optional[datetime] = None
 
 
 @dataclass