211 lines
7.7 KiB
Python
211 lines
7.7 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
import uuid
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, Any, List, Tuple
|
|
from dataclasses import dataclass, field, asdict
|
|
|
|
from config import settings
|
|
|
|
logger = logging.getLogger("AGYInviteManager")
|
|
|
|
@dataclass
|
|
class InviteToken:
|
|
code: str
|
|
creator_id: int
|
|
created_at: float = field(default_factory=time.time)
|
|
max_uses: int = 1 # 1 = single-use, 0 = unlimited, >1 = limited
|
|
uses_count: int = 0
|
|
used_by: List[int] = field(default_factory=list)
|
|
expires_at: Optional[float] = None
|
|
is_active: bool = True
|
|
note: str = ""
|
|
|
|
def is_valid(self) -> bool:
|
|
if not self.is_active:
|
|
return False
|
|
if self.expires_at and time.time() > self.expires_at:
|
|
return False
|
|
if self.max_uses > 0 and self.uses_count >= self.max_uses:
|
|
return False
|
|
return True
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> "InviteToken":
|
|
clean = {k: v for k, v in data.items() if k in cls.__dataclass_fields__}
|
|
return cls(**clean)
|
|
|
|
@dataclass
|
|
class PendingRequest:
|
|
user_id: int
|
|
username: str
|
|
first_name: str
|
|
last_name: str
|
|
requested_at: float = field(default_factory=time.time)
|
|
status: str = "pending" # "pending", "approved", "denied"
|
|
|
|
@property
|
|
def full_name(self) -> str:
|
|
parts = [self.first_name, self.last_name]
|
|
return " ".join([p for p in parts if p]).strip() or f"User {self.user_id}"
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> "PendingRequest":
|
|
clean = {k: v for k, v in data.items() if k in cls.__dataclass_fields__}
|
|
return cls(**clean)
|
|
|
|
class InviteManager:
|
|
def __init__(self, storage_file: Path = Path("/root/telegram-agy-bot/invites.json")):
|
|
self.storage_file = storage_file
|
|
self.invites: Dict[str, InviteToken] = {}
|
|
self.pending_requests: Dict[int, PendingRequest] = {}
|
|
self._load()
|
|
|
|
def _load(self):
|
|
if self.storage_file.exists():
|
|
try:
|
|
with open(self.storage_file, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
if isinstance(data, dict):
|
|
inv_data = data.get("invites", {})
|
|
for code, tdata in inv_data.items():
|
|
if isinstance(tdata, dict):
|
|
self.invites[code] = InviteToken.from_dict(tdata)
|
|
|
|
req_data = data.get("pending_requests", {})
|
|
for uid_str, rdata in req_data.items():
|
|
if isinstance(rdata, dict):
|
|
self.pending_requests[int(uid_str)] = PendingRequest.from_dict(rdata)
|
|
logger.info(f"Loaded {len(self.invites)} invite tokens and {len(self.pending_requests)} pending requests.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to load invites from {self.storage_file}: {e}")
|
|
|
|
def save(self):
|
|
try:
|
|
temp_file = self.storage_file.with_suffix(".tmp")
|
|
data = {
|
|
"invites": {code: t.to_dict() for code, t in self.invites.items()},
|
|
"pending_requests": {str(uid): r.to_dict() for uid, r in self.pending_requests.items()},
|
|
}
|
|
with open(temp_file, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2, ensure_ascii=False)
|
|
temp_file.replace(self.storage_file)
|
|
except Exception as e:
|
|
logger.error(f"Failed to save invites: {e}")
|
|
|
|
def create_invite(
|
|
self,
|
|
creator_id: int,
|
|
max_uses: int = 1,
|
|
duration_hours: Optional[float] = None,
|
|
note: str = "",
|
|
) -> InviteToken:
|
|
code = f"inv_{uuid.uuid4().hex[:8]}"
|
|
expires_at = (time.time() + duration_hours * 3600) if duration_hours else None
|
|
|
|
token = InviteToken(
|
|
code=code,
|
|
creator_id=creator_id,
|
|
created_at=time.time(),
|
|
max_uses=max_uses,
|
|
uses_count=0,
|
|
used_by=[],
|
|
expires_at=expires_at,
|
|
is_active=True,
|
|
note=note,
|
|
)
|
|
self.invites[code] = token
|
|
self.save()
|
|
logger.info(f"Created invite token {code} by admin {creator_id} (max_uses={max_uses})")
|
|
return token
|
|
|
|
def validate_and_use_invite(self, code: str, user_id: int) -> Tuple[bool, str, Optional[InviteToken]]:
|
|
token = self.invites.get(code)
|
|
if not token:
|
|
return False, "لینک دعوت نامعتبر است / Invalid invite token.", None
|
|
|
|
if not token.is_active:
|
|
return False, "این لینک دعوت غیرفعال شده است / Invite link is disabled.", token
|
|
|
|
if token.expires_at and time.time() > token.expires_at:
|
|
return False, "این لینک دعوت منقضی شده است / Invite link has expired.", token
|
|
|
|
if token.max_uses > 0 and token.uses_count >= token.max_uses:
|
|
return False, "ظرفیت استفاده از این لینک دعوت پر شده است / Invite link usage limit reached.", token
|
|
|
|
if user_id in token.used_by:
|
|
return True, "شما قبلاً با این لینک عضو شدهاید / Already joined with this invite.", token
|
|
|
|
# Consume usage
|
|
token.uses_count += 1
|
|
token.used_by.append(user_id)
|
|
if token.max_uses > 0 and token.uses_count >= token.max_uses:
|
|
token.is_active = False
|
|
|
|
# Authorize user
|
|
settings.add_authorized_user(user_id)
|
|
|
|
# Remove from pending requests if any
|
|
if user_id in self.pending_requests:
|
|
del self.pending_requests[user_id]
|
|
|
|
self.save()
|
|
logger.info(f"User {user_id} successfully used invite token {code} (created by {token.creator_id})")
|
|
return True, "عضویت با موفقیت انجام شد / Successfully joined!", token
|
|
|
|
def add_pending_request(self, user_id: int, username: str, first_name: str, last_name: str) -> PendingRequest:
|
|
req = PendingRequest(
|
|
user_id=user_id,
|
|
username=username or "",
|
|
first_name=first_name or "",
|
|
last_name=last_name or "",
|
|
requested_at=time.time(),
|
|
status="pending",
|
|
)
|
|
self.pending_requests[user_id] = req
|
|
self.save()
|
|
logger.info(f"New access request registered for user {user_id} (@{username})")
|
|
return req
|
|
|
|
def approve_request(self, user_id: int) -> bool:
|
|
settings.add_authorized_user(user_id)
|
|
if user_id in self.pending_requests:
|
|
self.pending_requests[user_id].status = "approved"
|
|
del self.pending_requests[user_id]
|
|
self.save()
|
|
logger.info(f"Access request approved for user {user_id}")
|
|
return True
|
|
|
|
def deny_request(self, user_id: int) -> bool:
|
|
if user_id in self.pending_requests:
|
|
self.pending_requests[user_id].status = "denied"
|
|
del self.pending_requests[user_id]
|
|
self.save()
|
|
logger.info(f"Access request denied for user {user_id}")
|
|
return True
|
|
|
|
def get_pending_requests(self) -> List[PendingRequest]:
|
|
return list(self.pending_requests.values())
|
|
|
|
def get_active_invites(self) -> List[InviteToken]:
|
|
return [t for t in self.invites.values() if t.is_valid()]
|
|
|
|
def revoke_invite(self, code: str) -> bool:
|
|
if code in self.invites:
|
|
self.invites[code].is_active = False
|
|
self.save()
|
|
return True
|
|
return False
|
|
|
|
# Singleton instance
|
|
invite_manager = InviteManager()
|