Yuriy Yuriev
329 lines
12 KiB
Python
329 lines
12 KiB
Python
"""
|
|
Сохранение и загрузка задач между перезапусками.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import Dict, Optional
|
|
|
|
from managers.task_manager import TaskParams
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TaskStorage:
|
|
"""Хранилище задач в JSON файле."""
|
|
|
|
def __init__(self, file_path: str = "data/tasks.json"):
|
|
self.file_path = Path(file_path)
|
|
self.file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = asyncio.Lock()
|
|
|
|
def _params_to_dict(self, params: TaskParams) -> dict:
|
|
return {
|
|
"url": params.url,
|
|
"task_type": params.task_type,
|
|
"min_delay": params.min_delay,
|
|
"max_delay": params.max_delay,
|
|
"min_reading": params.min_reading,
|
|
"max_reading": params.max_reading,
|
|
"max_visits": params.max_visits,
|
|
"current_visit": params.current_visit,
|
|
"channel": params.channel,
|
|
"target_username": params.target_username,
|
|
"visits_per_link": params.visits_per_link,
|
|
"monitor_minutes": params.monitor_minutes,
|
|
"allowed_domains": params.allowed_domains,
|
|
"paused": params.paused,
|
|
"completed": params.completed,
|
|
"stopped": params.stopped,
|
|
"total_visits": params.total_visits,
|
|
"successful_visits": params.successful_visits,
|
|
"links_found": params.links_found,
|
|
"chat_id": params.chat_id,
|
|
"user_id": params.user_id,
|
|
"visits_percent": params.visits_percent,
|
|
"started_at": params.started_at.isoformat() if params.started_at else None,
|
|
}
|
|
|
|
def _dict_to_params(self, data: dict) -> TaskParams:
|
|
params = TaskParams(
|
|
url=data.get("url", ""),
|
|
task_type=data.get("task_type", "visit"),
|
|
min_delay=data.get("min_delay", 10),
|
|
max_delay=data.get("max_delay", 30),
|
|
min_reading=data.get("min_reading", 5),
|
|
max_reading=data.get("max_reading", 15),
|
|
max_visits=data.get("max_visits"),
|
|
current_visit=data.get("current_visit", 0),
|
|
channel=data.get("channel", ""),
|
|
target_username=data.get("target_username", "*"),
|
|
visits_per_link=data.get("visits_per_link", 5),
|
|
monitor_minutes=data.get("monitor_minutes", 10),
|
|
allowed_domains=data.get("allowed_domains"),
|
|
paused=data.get("paused", False),
|
|
completed=data.get("completed", False),
|
|
stopped=data.get("stopped", False),
|
|
total_visits=data.get("total_visits", 0),
|
|
successful_visits=data.get("successful_visits", 0),
|
|
links_found=data.get("links_found", 0),
|
|
chat_id=data.get("chat_id"),
|
|
user_id=data.get("user_id"),
|
|
visits_percent=data.get("visits_percent", 0.0),
|
|
)
|
|
if data.get("started_at"):
|
|
try:
|
|
params.started_at = datetime.fromisoformat(data["started_at"])
|
|
except (ValueError, TypeError) as e:
|
|
logger.warning(f"Failed to parse started_at: {e}")
|
|
return params
|
|
|
|
def _save_sync(self, data: dict) -> None:
|
|
with open(self.file_path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
def _load_sync(self) -> dict:
|
|
if not self.file_path.exists():
|
|
return {}
|
|
try:
|
|
with open(self.file_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, IOError) as e:
|
|
logger.warning(f"Failed to load tasks: {e}")
|
|
return {}
|
|
|
|
async def save_tasks(self, tasks: Dict[str, TaskParams]) -> None:
|
|
async with self._lock:
|
|
try:
|
|
data = {tid: self._params_to_dict(p) for tid, p in tasks.items()}
|
|
await asyncio.to_thread(self._save_sync, data)
|
|
logger.info(f"💾 Saved {len(data)} tasks to {self.file_path}")
|
|
except Exception as e:
|
|
logger.error(f"Save error: {e}")
|
|
|
|
async def load_tasks(self) -> Dict[str, TaskParams]:
|
|
try:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
if not data:
|
|
logger.info("No saved tasks found")
|
|
return {}
|
|
tasks = {tid: self._dict_to_params(d) for tid, d in data.items()}
|
|
logger.info(f"📂 Loaded {len(tasks)} tasks from {self.file_path}")
|
|
return tasks
|
|
except Exception as e:
|
|
logger.error(f"Load error: {e}")
|
|
return {}
|
|
|
|
async def save_task(self, task_id: str, params: TaskParams) -> None:
|
|
tasks = await self.load_tasks()
|
|
tasks[task_id] = params
|
|
await self.save_tasks(tasks)
|
|
|
|
async def delete_task(self, task_id: str) -> None:
|
|
tasks = await self.load_tasks()
|
|
if task_id in tasks:
|
|
del tasks[task_id]
|
|
await self.save_tasks(tasks)
|
|
|
|
|
|
class ChatStorage:
|
|
"""Хранилище привязки стримеров к чатам."""
|
|
|
|
def __init__(self, file_path: str = "data/chats.json"):
|
|
self.file_path = Path(file_path)
|
|
self.file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = asyncio.Lock()
|
|
|
|
def _load_sync(self) -> dict:
|
|
if not self.file_path.exists():
|
|
return {}
|
|
try:
|
|
with open(self.file_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, IOError) as e:
|
|
logger.warning(f"Failed to load chats: {e}")
|
|
return {}
|
|
|
|
def _save_sync(self, data: dict) -> None:
|
|
with open(self.file_path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
async def get_chat_streamers(self, chat_id: int) -> list:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
return data.get(str(chat_id), [])
|
|
|
|
async def add_streamer(self, chat_id: int, channel: str) -> None:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
key = str(chat_id)
|
|
if key not in data:
|
|
data[key] = []
|
|
if channel not in data[key]:
|
|
data[key].append(channel)
|
|
await asyncio.to_thread(self._save_sync, data)
|
|
|
|
async def remove_streamer(self, chat_id: int, channel: str) -> None:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
key = str(chat_id)
|
|
if key in data and channel in data[key]:
|
|
data[key].remove(channel)
|
|
await asyncio.to_thread(self._save_sync, data)
|
|
|
|
|
|
class _IntBalanceStorage:
|
|
"""Базовое хранилище числового баланса пользователей."""
|
|
|
|
def __init__(self, file_path: str):
|
|
self.file_path = Path(file_path)
|
|
self.file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = asyncio.Lock()
|
|
|
|
def _load_sync(self) -> dict:
|
|
if not self.file_path.exists():
|
|
return {}
|
|
try:
|
|
with open(self.file_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, IOError) as e:
|
|
logger.warning(f"Failed to load {self.file_path}: {e}")
|
|
return {}
|
|
|
|
def _save_sync(self, data: dict) -> None:
|
|
with open(self.file_path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
async def get_balance(self, user_id: int) -> int:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
return data.get(str(user_id), 0)
|
|
|
|
async def add_balance(self, user_id: int, amount: int) -> int:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
data[str(user_id)] = data.get(str(user_id), 0) + amount
|
|
await asyncio.to_thread(self._save_sync, data)
|
|
return data[str(user_id)]
|
|
|
|
async def deduct(self, user_id: int, amount: int) -> int:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
new_bal = max(0, data.get(str(user_id), 0) - amount)
|
|
data[str(user_id)] = new_bal
|
|
await asyncio.to_thread(self._save_sync, data)
|
|
return new_bal
|
|
|
|
async def get_all(self) -> dict:
|
|
async with self._lock:
|
|
return await asyncio.to_thread(self._load_sync)
|
|
|
|
|
|
class ChatHistoryStorage:
|
|
"""Хранит ID сообщений бота для очистки чата при рестарте."""
|
|
|
|
def __init__(self, file_path: str = "data/chat_history.json"):
|
|
self.file_path = Path(file_path)
|
|
self.file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = asyncio.Lock()
|
|
|
|
def _load_sync(self) -> dict:
|
|
if not self.file_path.exists():
|
|
return {}
|
|
try:
|
|
with open(self.file_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, IOError):
|
|
return {}
|
|
|
|
def _save_sync(self, data: dict) -> None:
|
|
with open(self.file_path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f)
|
|
|
|
async def add(self, user_id: int, msg_id: int) -> None:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
key = str(user_id)
|
|
ids = data.get(key, [])
|
|
ids.append(msg_id)
|
|
data[key] = ids[-200:] # храним последние 200
|
|
await asyncio.to_thread(self._save_sync, data)
|
|
|
|
async def get(self, user_id: int) -> list:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
return data.get(str(user_id), [])
|
|
|
|
async def clear(self, user_id: int) -> None:
|
|
async with self._lock:
|
|
data = await asyncio.to_thread(self._load_sync)
|
|
data.pop(str(user_id), None)
|
|
await asyncio.to_thread(self._save_sync, data)
|
|
|
|
|
|
class BalanceStorage(_IntBalanceStorage):
|
|
"""Баланс переходов пользователей."""
|
|
|
|
def __init__(self, file_path: str = "data/balances.json"):
|
|
super().__init__(file_path)
|
|
|
|
|
|
class RubleBalanceStorage(_IntBalanceStorage):
|
|
"""Рублёвый баланс пользователей."""
|
|
|
|
def __init__(self, file_path: str = "data/rub_balances.json"):
|
|
super().__init__(file_path)
|
|
|
|
|
|
|
|
|
|
|
|
class PaymentStorage:
|
|
"""Хранилище ожидающих платежей."""
|
|
|
|
def __init__(self, file_path: str = "data/payments.json"):
|
|
self.file_path = Path(file_path)
|
|
self.file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = asyncio.Lock()
|
|
|
|
def _load_sync(self) -> dict:
|
|
if not self.file_path.exists():
|
|
return {}
|
|
try:
|
|
with open(self.file_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, IOError):
|
|
return {}
|
|
|
|
def _save_sync(self, data: dict) -> None:
|
|
with open(self.file_path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
async def save(self, payment_id: str, data: dict) -> None:
|
|
async with self._lock:
|
|
payments = await asyncio.to_thread(self._load_sync)
|
|
payments[payment_id] = data
|
|
await asyncio.to_thread(self._save_sync, payments)
|
|
|
|
async def get(self, payment_id: str) -> Optional[dict]:
|
|
async with self._lock:
|
|
payments = await asyncio.to_thread(self._load_sync)
|
|
return payments.get(payment_id)
|
|
|
|
async def delete(self, payment_id: str) -> None:
|
|
async with self._lock:
|
|
payments = await asyncio.to_thread(self._load_sync)
|
|
payments.pop(payment_id, None)
|
|
await asyncio.to_thread(self._save_sync, payments)
|
|
|
|
async def get_by_user(self, user_id: int) -> Optional[dict]:
|
|
async with self._lock:
|
|
payments = await asyncio.to_thread(self._load_sync)
|
|
for pid, data in payments.items():
|
|
if data.get("user_id") == user_id:
|
|
return {**data, "payment_id": pid}
|
|
return None
|