359 lines
12 KiB
Python
359 lines
12 KiB
Python
import threading
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
from urllib.parse import parse_qs, urlparse
|
|
|
|
import requests
|
|
|
|
from app.services.crypto_service import decrypt_value, encrypt_value
|
|
from app.services.db import db_transaction
|
|
from app.services.email_service import send_email_async
|
|
from app.services.zerodha_service import exchange_request_token
|
|
from app.services.zerodha_storage import set_session
|
|
from app.broker_store import expire_user_broker_session
|
|
|
|
IST = timezone(timedelta(hours=5, minutes=30))
|
|
|
|
KITE_LOGIN_ENDPOINT = "https://kite.zerodha.com/api/login"
|
|
KITE_TWOFA_ENDPOINT = "https://kite.zerodha.com/api/twofa"
|
|
|
|
|
|
class AutoLoginError(Exception):
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DB helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def save_auto_login_credentials(
|
|
user_id: str,
|
|
zerodha_login_id: str,
|
|
password: str,
|
|
totp_secret: str,
|
|
) -> None:
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE user_broker
|
|
SET zerodha_login_id = %s,
|
|
zerodha_password = %s,
|
|
totp_secret = %s,
|
|
auto_login_enabled = TRUE
|
|
WHERE user_id = %s
|
|
""",
|
|
(
|
|
zerodha_login_id.strip(),
|
|
encrypt_value(password),
|
|
encrypt_value(totp_secret.strip().replace(" ", "")),
|
|
user_id,
|
|
),
|
|
)
|
|
|
|
|
|
def get_auto_login_credentials(user_id: str) -> dict | None:
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT zerodha_login_id, zerodha_password, totp_secret,
|
|
auto_login_enabled, auto_login_last_at, auto_login_last_err,
|
|
api_key, api_secret
|
|
FROM user_broker
|
|
WHERE user_id = %s
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return None
|
|
login_id, enc_password, enc_totp, enabled, last_at, last_err, api_key, enc_api_secret = row
|
|
if not enabled or not login_id or not enc_password or not enc_totp:
|
|
return None
|
|
return {
|
|
"zerodha_login_id": login_id,
|
|
"password": decrypt_value(enc_password),
|
|
"totp_secret": decrypt_value(enc_totp),
|
|
"api_key": api_key,
|
|
"api_secret": decrypt_value(enc_api_secret) if enc_api_secret else None,
|
|
"last_refreshed_at": last_at.isoformat() if last_at else None,
|
|
"last_error": last_err,
|
|
}
|
|
|
|
|
|
def get_auto_login_status(user_id: str) -> dict:
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT auto_login_enabled, auto_login_last_at, auto_login_last_err,
|
|
zerodha_login_id
|
|
FROM user_broker
|
|
WHERE user_id = %s
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return {"configured": False}
|
|
enabled, last_at, last_err, login_id = row
|
|
return {
|
|
"configured": bool(enabled and login_id),
|
|
"last_refreshed_at": last_at.isoformat() if last_at else None,
|
|
"last_error": last_err,
|
|
}
|
|
|
|
|
|
def delete_auto_login_credentials(user_id: str) -> None:
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE user_broker
|
|
SET zerodha_login_id = NULL,
|
|
zerodha_password = NULL,
|
|
totp_secret = NULL,
|
|
auto_login_enabled = FALSE,
|
|
auto_login_last_at = NULL,
|
|
auto_login_last_err = NULL
|
|
WHERE user_id = %s
|
|
""",
|
|
(user_id,),
|
|
)
|
|
|
|
|
|
def _update_auto_login_result(user_id: str, error: str | None) -> None:
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE user_broker
|
|
SET auto_login_last_at = NOW(),
|
|
auto_login_last_err = %s
|
|
WHERE user_id = %s
|
|
""",
|
|
(error, user_id),
|
|
)
|
|
|
|
|
|
def _get_all_auto_login_users() -> list[dict]:
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT ub.user_id, au.username,
|
|
ub.zerodha_login_id, ub.zerodha_password, ub.totp_secret,
|
|
ub.api_key, ub.api_secret
|
|
FROM user_broker ub
|
|
JOIN app_user au ON au.id = ub.user_id
|
|
WHERE ub.auto_login_enabled = TRUE
|
|
AND ub.zerodha_login_id IS NOT NULL
|
|
AND ub.zerodha_password IS NOT NULL
|
|
AND ub.totp_secret IS NOT NULL
|
|
"""
|
|
)
|
|
rows = cur.fetchall()
|
|
results = []
|
|
for row in rows:
|
|
user_id, email, login_id, enc_pw, enc_totp, api_key, enc_secret = row
|
|
results.append({
|
|
"user_id": user_id,
|
|
"email": email,
|
|
"zerodha_login_id": login_id,
|
|
"password": decrypt_value(enc_pw),
|
|
"totp_secret": decrypt_value(enc_totp),
|
|
"api_key": api_key,
|
|
"api_secret": decrypt_value(enc_secret) if enc_secret else None,
|
|
})
|
|
return results
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core login flow
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _perform_zerodha_login(
|
|
zerodha_login_id: str,
|
|
password: str,
|
|
totp_secret: str,
|
|
api_key: str,
|
|
api_secret: str,
|
|
) -> dict:
|
|
"""Automates Zerodha login and returns session data with access_token."""
|
|
session = requests.Session()
|
|
session.headers.update({
|
|
"X-Kite-Version": "3",
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
|
|
})
|
|
|
|
# Step 1: Username + password
|
|
login_resp = session.post(
|
|
KITE_LOGIN_ENDPOINT,
|
|
data={"user_id": zerodha_login_id, "password": password},
|
|
timeout=15,
|
|
)
|
|
try:
|
|
login_data = login_resp.json()
|
|
except Exception:
|
|
raise AutoLoginError(f"Invalid response from Zerodha login: {login_resp.text[:200]}")
|
|
|
|
if login_data.get("status") != "success":
|
|
raise AutoLoginError(f"Zerodha login failed: {login_data.get('message', 'Unknown error')}")
|
|
|
|
request_id = login_data["data"]["request_id"]
|
|
|
|
# Step 2: TOTP — don't follow redirect automatically
|
|
try:
|
|
import pyotp
|
|
except ImportError:
|
|
raise AutoLoginError("pyotp is not installed on the server. Run: pip install pyotp==2.9.0")
|
|
totp_value = pyotp.TOTP(totp_secret).now()
|
|
twofa_resp = session.post(
|
|
KITE_TWOFA_ENDPOINT,
|
|
data={
|
|
"user_id": zerodha_login_id,
|
|
"request_id": request_id,
|
|
"twofa_value": totp_value,
|
|
"twofa_type": "totp",
|
|
},
|
|
timeout=15,
|
|
allow_redirects=False,
|
|
)
|
|
|
|
# Step 3: Follow redirects manually to intercept request_token
|
|
request_token = None
|
|
location = twofa_resp.headers.get("Location", "")
|
|
|
|
for _ in range(10):
|
|
if "request_token" in location:
|
|
parsed = urlparse(location)
|
|
params = parse_qs(parsed.query)
|
|
request_token = params.get("request_token", [None])[0]
|
|
break
|
|
if not location or twofa_resp.status_code not in (301, 302, 303, 307, 308):
|
|
break
|
|
twofa_resp = session.get(location, allow_redirects=False, timeout=15)
|
|
location = twofa_resp.headers.get("Location", "")
|
|
|
|
if not request_token:
|
|
raise AutoLoginError(
|
|
"Could not extract request_token from Zerodha redirect. "
|
|
"Check TOTP secret and credentials."
|
|
)
|
|
|
|
# Step 4: Exchange request_token for access_token using existing service
|
|
session_data = exchange_request_token(api_key, api_secret, request_token)
|
|
return {
|
|
"api_key": api_key,
|
|
"access_token": session_data.get("access_token"),
|
|
"request_token": request_token,
|
|
"user_name": session_data.get("user_name"),
|
|
"broker_user_id": session_data.get("user_id"),
|
|
}
|
|
|
|
|
|
def _reconnect_broker_after_auto_login(user_id: str) -> None:
|
|
"""Marks broker as connected after successful auto-login."""
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE user_broker
|
|
SET connected = TRUE,
|
|
auth_state = 'CONNECTED'
|
|
WHERE user_id = %s
|
|
""",
|
|
(user_id,),
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public: execute for one user
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def execute_auto_login(user_id: str, email: str | None = None) -> dict:
|
|
"""Run auto-login for a single user. Returns result dict."""
|
|
creds = get_auto_login_credentials(user_id)
|
|
if not creds:
|
|
return {"success": False, "error": "Auto-login not configured"}
|
|
|
|
try:
|
|
session_data = _perform_zerodha_login(
|
|
zerodha_login_id=creds["zerodha_login_id"],
|
|
password=creds["password"],
|
|
totp_secret=creds["totp_secret"],
|
|
api_key=creds["api_key"],
|
|
api_secret=creds["api_secret"],
|
|
)
|
|
set_session(user_id, session_data)
|
|
_reconnect_broker_after_auto_login(user_id)
|
|
_update_auto_login_result(user_id, error=None)
|
|
if email:
|
|
send_email_async(
|
|
email,
|
|
"Zerodha session refreshed automatically",
|
|
(
|
|
"Your Zerodha session has been refreshed automatically by QuantFortune.\n\n"
|
|
"Your strategy will continue running without any interruption.\n\n"
|
|
f"Refreshed at: {datetime.now(IST).strftime('%d %b %Y, %I:%M %p IST')}"
|
|
),
|
|
)
|
|
print(f"[AUTO-LOGIN] Successfully refreshed session for user {user_id}", flush=True)
|
|
return {"success": True}
|
|
except Exception as exc:
|
|
err_msg = str(exc)
|
|
_update_auto_login_result(user_id, error=err_msg)
|
|
if email:
|
|
send_email_async(
|
|
email,
|
|
"Action required: Zerodha auto-login failed",
|
|
(
|
|
f"QuantFortune could not automatically refresh your Zerodha session.\n\n"
|
|
f"Error: {err_msg}\n\n"
|
|
"Please log in to QuantFortune and reconnect your Zerodha account manually.\n\n"
|
|
"If your strategy was running, it has been paused until you reconnect."
|
|
),
|
|
)
|
|
print(f"[AUTO-LOGIN] Failed for user {user_id}: {exc}", flush=True)
|
|
return {"success": False, "error": err_msg}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Refresh all users (called by daily scheduler)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def refresh_all_auto_login_sessions() -> None:
|
|
users = _get_all_auto_login_users()
|
|
print(f"[AUTO-LOGIN] Starting daily refresh for {len(users)} user(s)", flush=True)
|
|
for user in users:
|
|
execute_auto_login(user_id=user["user_id"], email=user["email"])
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Daily scheduler — runs at 6:05 AM IST every day
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _next_refresh_time() -> datetime:
|
|
now = datetime.now(IST)
|
|
target = now.replace(hour=6, minute=5, second=0, microsecond=0)
|
|
if now >= target:
|
|
target += timedelta(days=1)
|
|
return target
|
|
|
|
|
|
def _scheduler_loop() -> None:
|
|
while True:
|
|
next_run = _next_refresh_time()
|
|
sleep_seconds = (_next_refresh_time() - datetime.now(IST)).total_seconds()
|
|
print(
|
|
f"[AUTO-LOGIN] Next scheduled refresh at "
|
|
f"{next_run.strftime('%d %b %Y %I:%M %p IST')}, "
|
|
f"sleeping {sleep_seconds:.0f}s",
|
|
flush=True,
|
|
)
|
|
time.sleep(sleep_seconds)
|
|
try:
|
|
refresh_all_auto_login_sessions()
|
|
except Exception as exc:
|
|
print(f"[AUTO-LOGIN] Scheduler error: {exc}", flush=True)
|
|
|
|
|
|
def start_auto_login_scheduler() -> None:
|
|
thread = threading.Thread(target=_scheduler_loop, daemon=True, name="auto-login-scheduler")
|
|
thread.start()
|
|
print("[AUTO-LOGIN] Daily scheduler started (fires at 6:05 AM IST)", flush=True)
|