Kalshi-Bot/kalshi_client.py

from __future__ import annotations

import base64
import dataclasses
import datetime as dt
import os
import time
from typing import Any, Dict, Optional

import requests
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa


@dataclasses.dataclass(frozen=True)
class KalshiConfig:
    env: str  # "demo" or "prod"
    api_key_id: str
    private_key_path: str

    @property
    def base_url(self) -> str:
        # Docs show demo at demo-api.kalshi.co and public market data at api.elections.kalshi.com.
        # We’ll use:
        # - demo trading: https://demo-api.kalshi.co
        # - prod/public: https://api.elections.kalshi.com
        if self.env.lower() == "demo":
            return "https://demo-api.kalshi.co"
        return "https://api.elections.kalshi.com"


def _load_private_key_from_file(file_path: str) -> rsa.RSAPrivateKey:
    with open(file_path, "rb") as key_file:
        private_key = serialization.load_pem_private_key(
            key_file.read(),
            password=None,
            backend=default_backend(),
        )
    if not isinstance(private_key, rsa.RSAPrivateKey):
        raise ValueError("Private key is not an RSA private key")
    return private_key


def _sign_pss_text(private_key: rsa.RSAPrivateKey, text: str) -> str:
    message = text.encode("utf-8")
    signature = private_key.sign(
        message,
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.DIGEST_LENGTH,
        ),
        hashes.SHA256(),
    )
    return base64.b64encode(signature).decode("utf-8")


class KalshiClient:
    """
    Minimal REST client using Kalshi's signed headers:
    KALSHI-ACCESS-KEY, KALSHI-ACCESS-TIMESTAMP (ms), KALSHI-ACCESS-SIGNATURE
    Signature = RSA-PSS(SHA256) over: timestamp + METHOD + path_without_query
    """
    def __init__(self, cfg: KalshiConfig, session: Optional[requests.Session] = None):
        self.cfg = cfg
        self.session = session or requests.Session()
        self._private_key = _load_private_key_from_file(cfg.private_key_path)

    @staticmethod
    def from_env() -> "KalshiClient":
        env = os.getenv("KALSHI_ENV", "demo")
        api_key_id = os.environ["KALSHI_API_KEY_ID"]
        key_path = os.environ["KALSHI_PRIVATE_KEY_PATH"]
        return KalshiClient(KalshiConfig(env=env, api_key_id=api_key_id, private_key_path=key_path))

    def _auth_headers(self, method: str, path: str) -> Dict[str, str]:
        ts = str(int(time.time() * 1000))
        path_wo_query = path.split("?")[0]
        msg = f"{ts}{method.upper()}{path_wo_query}"
        sig = _sign_pss_text(self._private_key, msg)
        return {
            "KALSHI-ACCESS-KEY": self.cfg.api_key_id,
            "KALSHI-ACCESS-TIMESTAMP": ts,
            "KALSHI-ACCESS-SIGNATURE": sig,
        }

    def _request(self, method: str, path: str, *, params: Optional[dict] = None, json: Optional[dict] = None, auth: bool = False) -> Dict[str, Any]:
        url = self.cfg.base_url + path
        headers: Dict[str, str] = {"Content-Type": "application/json"}
        if auth:
            headers.update(self._auth_headers(method, path))

        resp = self.session.request(method=method, url=url, params=params, json=json, headers=headers, timeout=15)
        if resp.status_code >= 400:
            raise RuntimeError(f"Kalshi API error {resp.status_code}: {resp.text}")
        if resp.status_code == 204:
            return {}
        return resp.json()

    # ---- Public market-data endpoints (no auth required in docs) ----
    def get_series_list(self, *, category: str) -> Dict[str, Any]:
        return self._request("GET", "/trade-api/v2/series", params={"category": category}, auth=False)

    def get_markets(self, *, series_ticker: str, status: str = "open", limit: int = 200) -> Dict[str, Any]:
        params = {"series_ticker": series_ticker, "status": status, "limit": limit}
        return self._request("GET", "/trade-api/v2/markets", params=params, auth=False)

    def get_market(self, ticker: str) -> Dict[str, Any]:
        return self._request("GET", f"/trade-api/v2/markets/{ticker}", auth=False)

    # ---- Trading endpoints (auth required) ----
    def create_order(
        self,
        *,
        ticker: str,
        side: str,
        action: str,
        count: int,
        yes_price_cents: Optional[int] = None,
        no_price_cents: Optional[int] = None,
        buy_max_cost_cents: Optional[int] = None,
        time_in_force: str = "fill_or_kill",
        client_order_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "ticker": ticker,
            "side": side,           # "yes" or "no"
            "action": action,       # "buy" or "sell"
            "count": int(count),
            "type": "limit",
            "time_in_force": time_in_force,
        }
        if client_order_id:
            payload["client_order_id"] = client_order_id
        if yes_price_cents is not None:
            payload["yes_price"] = int(yes_price_cents)
        if no_price_cents is not None:
            payload["no_price"] = int(no_price_cents)
        if buy_max_cost_cents is not None:
            payload["buy_max_cost"] = int(buy_max_cost_cents)

        return self._request("POST", "/trade-api/v2/portfolio/orders", json=payload, auth=True)

    def get_order(self, order_id: str) -> Dict[str, Any]:
        return self._request("GET", f"/trade-api/v2/portfolio/orders/{order_id}", auth=True)

    def cancel_order(self, order_id: str) -> Dict[str, Any]:
        return self._request("DELETE", f"/trade-api/v2/portfolio/orders/{order_id}", auth=True)