2026-02-01 13:57:30 +00:00

206 lines
7.0 KiB
Python

import os
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import RedirectResponse
from app.broker_store import (
clear_user_broker,
get_broker_credentials,
get_pending_broker,
get_user_broker,
set_broker_auth_state,
set_connected_broker,
set_pending_broker,
)
from app.services.auth_service import get_user_for_session
from app.services.zerodha_service import build_login_url, exchange_request_token
from app.services.email_service import send_email
from app.services.zerodha_storage import set_session
router = APIRouter(prefix="/api/broker")
def _require_user(request: Request):
session_id = request.cookies.get("session_id")
if not session_id:
raise HTTPException(status_code=401, detail="Not authenticated")
user = get_user_for_session(session_id)
if not user:
raise HTTPException(status_code=401, detail="Not authenticated")
return user
@router.post("/connect")
async def connect_broker(payload: dict, request: Request):
user = _require_user(request)
broker = (payload.get("broker") or "").strip()
token = (payload.get("token") or "").strip()
user_name = (payload.get("userName") or "").strip()
broker_user_id = (payload.get("brokerUserId") or "").strip()
if not broker or not token:
raise HTTPException(status_code=400, detail="Broker and token are required")
set_connected_broker(
user["id"],
broker,
token,
user_name=user_name or None,
broker_user_id=broker_user_id or None,
)
try:
body = (
"Your broker has been connected to Quantfortune.\n\n"
f"Broker: {broker}\n"
f"Broker User ID: {broker_user_id or 'N/A'}\n"
)
send_email(user["username"], "Broker connected", body)
except Exception:
pass
return {"connected": True}
@router.get("/status")
async def broker_status(request: Request):
user = _require_user(request)
entry = get_user_broker(user["id"])
if not entry or not entry.get("connected"):
return {"connected": False}
return {
"connected": True,
"broker": entry.get("broker"),
"connected_at": entry.get("connected_at"),
"userName": entry.get("user_name"),
"brokerUserId": entry.get("broker_user_id"),
"authState": entry.get("auth_state"),
}
@router.post("/disconnect")
async def disconnect_broker(request: Request):
user = _require_user(request)
clear_user_broker(user["id"])
set_broker_auth_state(user["id"], "DISCONNECTED")
try:
body = "Your broker connection has been disconnected from Quantfortune."
send_email(user["username"], "Broker disconnected", body)
except Exception:
pass
return {"connected": False}
@router.post("/zerodha/login")
async def zerodha_login(payload: dict, request: Request):
user = _require_user(request)
api_key = (payload.get("apiKey") or "").strip()
api_secret = (payload.get("apiSecret") or "").strip()
redirect_url = (payload.get("redirectUrl") or "").strip()
if not api_key or not api_secret:
raise HTTPException(status_code=400, detail="API key and secret are required")
set_pending_broker(user["id"], "ZERODHA", api_key, api_secret)
return {"loginUrl": build_login_url(api_key, redirect_url=redirect_url or None)}
@router.get("/zerodha/callback")
async def zerodha_callback(request: Request, request_token: str = ""):
user = _require_user(request)
token = request_token.strip()
if not token:
raise HTTPException(status_code=400, detail="Missing request_token")
pending = get_pending_broker(user["id"]) or {}
api_key = (pending.get("api_key") or "").strip()
api_secret = (pending.get("api_secret") or "").strip()
if not api_key or not api_secret:
raise HTTPException(status_code=400, detail="Zerodha login not initialized")
try:
session_data = exchange_request_token(api_key, api_secret, token)
except Exception as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
access_token = session_data.get("access_token")
if not access_token:
raise HTTPException(status_code=400, detail="Missing access token from Zerodha")
saved = set_session(
user["id"],
{
"api_key": api_key,
"access_token": access_token,
"request_token": session_data.get("request_token", token),
"user_name": session_data.get("user_name"),
"broker_user_id": session_data.get("user_id"),
},
)
set_connected_broker(
user["id"],
"ZERODHA",
access_token,
api_key=api_key,
api_secret=api_secret,
user_name=session_data.get("user_name"),
broker_user_id=session_data.get("user_id"),
auth_state="VALID",
)
return {
"connected": True,
"userName": saved.get("user_name"),
"brokerUserId": saved.get("broker_user_id"),
}
@router.get("/login")
async def broker_login(request: Request):
user = _require_user(request)
creds = get_broker_credentials(user["id"])
if not creds:
raise HTTPException(status_code=400, detail="Broker credentials not configured")
redirect_url = (os.getenv("ZERODHA_REDIRECT_URL") or "").strip()
if not redirect_url:
base = str(request.base_url).rstrip("/")
redirect_url = f"{base}/api/broker/callback"
login_url = build_login_url(creds["api_key"], redirect_url=redirect_url)
return RedirectResponse(login_url)
@router.get("/callback")
async def broker_callback(request: Request, request_token: str = ""):
user = _require_user(request)
token = request_token.strip()
if not token:
raise HTTPException(status_code=400, detail="Missing request_token")
creds = get_broker_credentials(user["id"])
if not creds:
raise HTTPException(status_code=400, detail="Broker credentials not configured")
try:
session_data = exchange_request_token(creds["api_key"], creds["api_secret"], token)
except Exception as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
access_token = session_data.get("access_token")
if not access_token:
raise HTTPException(status_code=400, detail="Missing access token from Zerodha")
set_session(
user["id"],
{
"api_key": creds["api_key"],
"access_token": access_token,
"request_token": session_data.get("request_token", token),
"user_name": session_data.get("user_name"),
"broker_user_id": session_data.get("user_id"),
},
)
set_connected_broker(
user["id"],
"ZERODHA",
access_token,
api_key=creds["api_key"],
api_secret=creds["api_secret"],
user_name=session_data.get("user_name"),
broker_user_id=session_data.get("user_id"),
auth_state="VALID",
)
target_url = os.getenv("BROKER_DASHBOARD_URL") or "/dashboard?armed=false"
return RedirectResponse(target_url)