445 lines
14 KiB
Python
445 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
from datetime import date, datetime, timedelta, timezone
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
|
|
from bot.db.repository import Budget, Expense
|
|
from bot.services.budget import (
|
|
PERIOD_PAGE_SIZE,
|
|
BudgetStatus,
|
|
PeriodExpensesPage,
|
|
)
|
|
|
|
|
|
YANDEX_NAMESPACE_BIT = 1 << 50
|
|
|
|
|
|
class BudgetApiError(Exception):
|
|
"""HTTP / transport failure talking to the API."""
|
|
|
|
|
|
class YandexLoginRequired(Exception):
|
|
"""Bot identity is not linked to a Yandex account yet."""
|
|
|
|
def __init__(self, login_url: str, detail: str = "") -> None:
|
|
self.login_url = login_url
|
|
self.detail = detail or "Чтобы пользоваться ботом, войдите через Яндекс."
|
|
super().__init__(self.detail)
|
|
|
|
|
|
def _jwt_uid(token: str) -> int | None:
|
|
try:
|
|
payload = token.split(".")[1]
|
|
payload += "=" * (-len(payload) % 4)
|
|
data = json.loads(base64.urlsafe_b64decode(payload.encode("ascii")))
|
|
raw = data.get("uid") or data.get("sub")
|
|
return int(raw)
|
|
except (IndexError, ValueError, TypeError, json.JSONDecodeError, OSError):
|
|
return None
|
|
|
|
|
|
class BudgetApiClient:
|
|
"""Budget operations via the same HTTP API as the web cabinet."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
base_url: str,
|
|
api_token: str,
|
|
session: aiohttp.ClientSession | None = None,
|
|
) -> None:
|
|
self._base = base_url.rstrip("/")
|
|
self._api_token = api_token
|
|
self._session = session
|
|
self._owns_session = session is None
|
|
self._tokens: dict[int, tuple[str, datetime]] = {}
|
|
|
|
async def start(self) -> None:
|
|
if self._session is None:
|
|
self._session = aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=30),
|
|
)
|
|
|
|
async def close(self) -> None:
|
|
if self._owns_session and self._session is not None:
|
|
await self._session.close()
|
|
self._session = None
|
|
|
|
def _ensure_session(self) -> aiohttp.ClientSession:
|
|
if self._session is None:
|
|
raise RuntimeError("BudgetApiClient is not started")
|
|
return self._session
|
|
|
|
@staticmethod
|
|
def _detail(payload: Any, fallback: str) -> str:
|
|
if isinstance(payload, dict):
|
|
detail = payload.get("detail", fallback)
|
|
if isinstance(detail, str):
|
|
return detail
|
|
if isinstance(detail, list) and detail:
|
|
first = detail[0]
|
|
if isinstance(first, dict) and "msg" in first:
|
|
return str(first["msg"])
|
|
return str(first)
|
|
return str(detail)
|
|
return fallback
|
|
|
|
async def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
user_id: int | None = None,
|
|
json_body: dict[str, Any] | None = None,
|
|
params: dict[str, Any] | None = None,
|
|
retry_auth: bool = True,
|
|
) -> Any:
|
|
session = self._ensure_session()
|
|
headers: dict[str, str] = {"Accept": "application/json"}
|
|
if user_id is not None:
|
|
headers["Authorization"] = f"Bearer {await self._access_token(user_id)}"
|
|
else:
|
|
headers["X-API-Token"] = self._api_token
|
|
|
|
clean_params = None
|
|
if params:
|
|
clean_params = {k: v for k, v in params.items() if v is not None}
|
|
|
|
url = f"{self._base}{path}"
|
|
async with session.request(
|
|
method,
|
|
url,
|
|
json=json_body,
|
|
params=clean_params,
|
|
headers=headers,
|
|
) as resp:
|
|
if resp.status == 401 and user_id is not None and retry_auth:
|
|
self._tokens.pop(user_id, None)
|
|
await resp.read()
|
|
return await self._request(
|
|
method,
|
|
path,
|
|
user_id=user_id,
|
|
json_body=json_body,
|
|
params=params,
|
|
retry_auth=False,
|
|
)
|
|
|
|
if resp.status == 204:
|
|
return None
|
|
|
|
body: Any
|
|
try:
|
|
body = await resp.json(content_type=None)
|
|
except aiohttp.ContentTypeError:
|
|
text = await resp.text()
|
|
body = {"detail": text or resp.reason}
|
|
|
|
if resp.status >= 400:
|
|
if (
|
|
resp.status == 403
|
|
and isinstance(body, dict)
|
|
and body.get("code") == "yandex_required"
|
|
):
|
|
if user_id is not None:
|
|
self._tokens.pop(user_id, None)
|
|
raise YandexLoginRequired(
|
|
str(body.get("login_url") or ""),
|
|
self._detail(body, "Чтобы пользоваться ботом, войдите через Яндекс."),
|
|
)
|
|
message = self._detail(body, f"API error {resp.status}")
|
|
if resp.status in {400, 404}:
|
|
raise ValueError(message)
|
|
raise BudgetApiError(message)
|
|
return body
|
|
|
|
async def ensure_yandex_login(self, telegram_user_id: int) -> None:
|
|
await self._access_token(telegram_user_id)
|
|
|
|
async def _access_token(self, user_id: int) -> str:
|
|
now = datetime.now(timezone.utc)
|
|
cached = self._tokens.get(user_id)
|
|
if cached and cached[1] > now + timedelta(seconds=60):
|
|
uid = _jwt_uid(cached[0])
|
|
if uid is not None and (uid & YANDEX_NAMESPACE_BIT):
|
|
return cached[0]
|
|
self._tokens.pop(user_id, None)
|
|
|
|
data = await self._request(
|
|
"POST",
|
|
"/api/auth/internal",
|
|
json_body={"user_id": user_id},
|
|
)
|
|
token = str(data["access_token"])
|
|
uid = _jwt_uid(token)
|
|
if uid is None or not (uid & YANDEX_NAMESPACE_BIT):
|
|
raise BudgetApiError("API did not issue a Yandex-linked session")
|
|
self._tokens[user_id] = (token, now + timedelta(days=1))
|
|
return token
|
|
|
|
@staticmethod
|
|
def _parse_date(value: str | date) -> date:
|
|
if isinstance(value, date) and not isinstance(value, datetime):
|
|
return value
|
|
return date.fromisoformat(str(value)[:10])
|
|
|
|
def _status_from(self, data: dict[str, Any]) -> BudgetStatus:
|
|
b = data["budget"]
|
|
budget = Budget(
|
|
id=int(b["id"]),
|
|
user_id=int(b["user_id"]),
|
|
name=str(b.get("name") or "Бюджет"),
|
|
total_amount=float(b["total_amount"]),
|
|
start_date=self._parse_date(b["start_date"]),
|
|
end_date=self._parse_date(b["end_date"]),
|
|
currency=str(b.get("currency") or "RUB"),
|
|
is_active=bool(b.get("is_active", True)),
|
|
)
|
|
return BudgetStatus(
|
|
budget=budget,
|
|
today=self._parse_date(data["today"]),
|
|
days_left=int(data["days_left"]),
|
|
total_spent=float(data["total_spent"]),
|
|
remaining=float(data["remaining"]),
|
|
daily_limit=float(data["daily_limit"]),
|
|
spent_today=float(data["spent_today"]),
|
|
remaining_today=float(data["remaining_today"]),
|
|
is_over_daily=bool(data["is_over_daily"]),
|
|
is_over_budget=bool(data["is_over_budget"]),
|
|
is_expired=bool(data["is_expired"]),
|
|
selected=bool(data.get("selected", False)),
|
|
)
|
|
|
|
def _expense_from(self, data: dict[str, Any], user_id: int) -> Expense:
|
|
return Expense(
|
|
id=int(data["id"]),
|
|
user_id=user_id,
|
|
budget_id=int(data.get("budget_id") or 0),
|
|
amount=float(data["amount"]),
|
|
note=data.get("note"),
|
|
spent_at=self._parse_date(data["spent_at"]),
|
|
)
|
|
|
|
async def _page_from(
|
|
self,
|
|
user_id: int,
|
|
data: dict[str, Any],
|
|
budget_id: int | None = None,
|
|
) -> PeriodExpensesPage:
|
|
status = await self.get_status(
|
|
user_id,
|
|
budget_id=budget_id or int(data.get("budget_id") or 0) or None,
|
|
)
|
|
return PeriodExpensesPage(
|
|
budget=status.budget,
|
|
page=int(data["page"]),
|
|
total_pages=int(data["total_pages"]),
|
|
total_count=int(data["total_count"]),
|
|
total_sum=float(data["total_sum"]),
|
|
page_size=int(data["page_size"]),
|
|
items=[self._expense_from(item, user_id) for item in data.get("items", [])],
|
|
)
|
|
|
|
async def list_user_statuses(self, user_id: int) -> list[BudgetStatus]:
|
|
data = await self._request("GET", "/api/me/budgets", user_id=user_id)
|
|
return [self._status_from(item) for item in data.get("items", [])]
|
|
|
|
async def get_status(
|
|
self,
|
|
user_id: int,
|
|
today: date | None = None,
|
|
budget_id: int | None = None,
|
|
) -> BudgetStatus:
|
|
del today
|
|
data = await self._request(
|
|
"GET",
|
|
"/api/me/budget",
|
|
user_id=user_id,
|
|
params={"budget_id": budget_id} if budget_id else None,
|
|
)
|
|
return self._status_from(data)
|
|
|
|
async def create_budget(
|
|
self,
|
|
user_id: int,
|
|
total_amount: float,
|
|
end_date: date,
|
|
*,
|
|
name: str = "Бюджет",
|
|
start_date: date | None = None,
|
|
is_active: bool = True,
|
|
select: bool = True,
|
|
) -> BudgetStatus:
|
|
body: dict[str, Any] = {
|
|
"total_amount": total_amount,
|
|
"end_date": end_date.isoformat(),
|
|
"name": name,
|
|
"is_active": is_active,
|
|
"select": select,
|
|
}
|
|
if start_date is not None:
|
|
body["start_date"] = start_date.isoformat()
|
|
data = await self._request(
|
|
"POST",
|
|
"/api/me/budgets",
|
|
user_id=user_id,
|
|
json_body=body,
|
|
)
|
|
return self._status_from(data)
|
|
|
|
async def set_budget_active(
|
|
self,
|
|
user_id: int,
|
|
budget_id: int,
|
|
is_active: bool,
|
|
) -> BudgetStatus:
|
|
data = await self._request(
|
|
"PATCH",
|
|
f"/api/me/budgets/{budget_id}/active",
|
|
user_id=user_id,
|
|
json_body={"is_active": is_active},
|
|
)
|
|
return self._status_from(data)
|
|
|
|
async def select_budget(self, user_id: int, budget_id: int) -> BudgetStatus:
|
|
data = await self._request(
|
|
"POST",
|
|
f"/api/me/budgets/{budget_id}/select",
|
|
user_id=user_id,
|
|
)
|
|
return self._status_from(data)
|
|
|
|
async def delete_budget(self, user_id: int, budget_id: int) -> None:
|
|
await self._request(
|
|
"DELETE",
|
|
f"/api/me/budgets/{budget_id}",
|
|
user_id=user_id,
|
|
)
|
|
|
|
async def set_budget(
|
|
self,
|
|
user_id: int,
|
|
total_amount: float,
|
|
end_date: date,
|
|
start_date: date | None = None,
|
|
reset_expenses: bool = True,
|
|
name: str | None = None,
|
|
budget_id: int | None = None,
|
|
) -> BudgetStatus:
|
|
# Creating a new named budget is the common bot flow.
|
|
if budget_id is None and name is not None:
|
|
return await self.create_budget(
|
|
user_id,
|
|
total_amount,
|
|
end_date,
|
|
name=name,
|
|
start_date=start_date,
|
|
)
|
|
body: dict[str, Any] = {
|
|
"total_amount": total_amount,
|
|
"end_date": end_date.isoformat(),
|
|
"reset_expenses": reset_expenses,
|
|
}
|
|
if name is not None:
|
|
body["name"] = name
|
|
if budget_id is not None:
|
|
body["budget_id"] = budget_id
|
|
if start_date is not None:
|
|
body["start_date"] = start_date.isoformat()
|
|
data = await self._request(
|
|
"PUT",
|
|
"/api/me/budget",
|
|
user_id=user_id,
|
|
json_body=body,
|
|
)
|
|
return self._status_from(data)
|
|
|
|
async def add_expense(
|
|
self,
|
|
user_id: int,
|
|
amount: float,
|
|
note: str | None = None,
|
|
spent_at: date | None = None,
|
|
budget_id: int | None = None,
|
|
) -> BudgetStatus:
|
|
body: dict[str, Any] = {"amount": amount}
|
|
if note is not None:
|
|
body["note"] = note
|
|
if spent_at is not None:
|
|
body["spent_at"] = spent_at.isoformat()
|
|
if budget_id is not None:
|
|
body["budget_id"] = budget_id
|
|
data = await self._request(
|
|
"POST",
|
|
"/api/me/expenses",
|
|
user_id=user_id,
|
|
json_body=body,
|
|
)
|
|
return self._status_from(data)
|
|
|
|
async def undo_last_expense(
|
|
self,
|
|
user_id: int,
|
|
budget_id: int | None = None,
|
|
) -> tuple[BudgetStatus | None, float | None]:
|
|
try:
|
|
data = await self._request(
|
|
"DELETE",
|
|
"/api/me/expenses/last",
|
|
user_id=user_id,
|
|
params={"budget_id": budget_id} if budget_id else None,
|
|
)
|
|
except ValueError:
|
|
return None, None
|
|
return self._status_from(data["status"]), float(data["deleted_amount"])
|
|
|
|
async def get_period_expenses_page(
|
|
self,
|
|
user_id: int,
|
|
page: int = 0,
|
|
page_size: int = PERIOD_PAGE_SIZE,
|
|
budget_id: int | None = None,
|
|
) -> PeriodExpensesPage:
|
|
data = await self._request(
|
|
"GET",
|
|
"/api/me/expenses",
|
|
user_id=user_id,
|
|
params={"page": page, "page_size": page_size, "budget_id": budget_id},
|
|
)
|
|
return await self._page_from(user_id, data, budget_id=budget_id)
|
|
|
|
async def expenses_on_date(
|
|
self,
|
|
user_id: int,
|
|
day: date,
|
|
budget_id: int | None = None,
|
|
) -> list[Expense]:
|
|
data = await self._request(
|
|
"GET",
|
|
"/api/me/expenses",
|
|
user_id=user_id,
|
|
params={
|
|
"spent_at": day.isoformat(),
|
|
"page": 0,
|
|
"page_size": 100,
|
|
"budget_id": budget_id,
|
|
},
|
|
)
|
|
return [self._expense_from(item, user_id) for item in data.get("items", [])]
|
|
|
|
async def today_expenses(
|
|
self,
|
|
user_id: int,
|
|
today: date | None = None,
|
|
budget_id: int | None = None,
|
|
) -> list[Expense]:
|
|
return await self.expenses_on_date(
|
|
user_id,
|
|
today or date.today(),
|
|
budget_id=budget_id,
|
|
)
|