Authentication

# pip install eth-account websockets
import asyncio
import json
import os

import websockets
from eth_account import Account
from eth_account.messages import encode_defunct


AUTH_DOMAIN = b"longshot:mm:ws-auth:v1:longshot.xyz"


async def authenticate(url: str) -> str:
    private_key = os.environ["MM_PRIVATE_KEY"]
    wallet_address = Account.from_key(private_key).address

    async with websockets.connect(url) as ws:
        await ws.send(json.dumps({"type": "auth"}))

        challenge = json.loads(await ws.recv())
        if challenge.get("type") != "auth_challenge":
            raise RuntimeError(f"unexpected message: {challenge}")

        nonce_hex = challenge["nonce"]
        timestamp = int(challenge["timestamp"])

        nonce = bytes.fromhex(nonce_hex.removeprefix("0x"))
        if len(nonce) != 32:
            raise ValueError("auth nonce must decode to 32 bytes")

        auth_message = AUTH_DOMAIN + nonce + timestamp.to_bytes(8, "little", signed=False)
        if len(auth_message) != 75:
            raise ValueError("auth message must be exactly 75 bytes")

        signed = Account.sign_message(
            encode_defunct(primitive=auth_message),
            private_key=private_key,
        )

        auth_response = {
            "type": "auth_response",
            "wallet_address": wallet_address,
            "signature": signed.signature.hex().removeprefix("0x"),
        }
        await ws.send(json.dumps(auth_response))

        result = json.loads(await ws.recv())
        if result.get("type") != "auth_result" or not result.get("success"):
            raise RuntimeError(f"auth failed: {result}")

        return result["session_token"]


asyncio.run(authenticate("ws://127.0.0.1:3001/ws"))