| from pyrogram import Client |
| from telethon.sync import TelegramClient |
| from telethon.sessions import StringSession |
| from typing import List, Optional, Dict |
| import os |
| import aiohttp |
| from pydantic import BaseModel |
|
|
|
|
| class Node(BaseModel): |
| MASTER: bool |
| SPACE_HOST: str |
|
|
|
|
| class NodeTaskState(BaseModel): |
| NODE: Node |
| CHUNK: int |
| COMPLETED: bool |
|
|
|
|
| class Task(BaseModel): |
| TASK_ID: str |
| COMPLETED: bool = False |
| NODES: Optional[List[NodeTaskState]] = [] |
|
|
| def is_completed(self) -> bool: |
| for node in self.NODES: |
| if not node.COMPLETED: |
| return False |
| self.COMPLETED = True |
| return True |
|
|
| def mark_node_completed(self, space_host: str): |
| for state in self.NODES: |
| if state.NODE.SPACE_HOST == space_host: |
| state.COMPLETED = True |
| break |
|
|
| async def add_node(self, node: Node, CHUNK: int): |
| new_node_state = NodeTaskState(NODE=node, CHUNK=CHUNK, COMPLETED=False) |
| self.NODES.append(new_node_state) |
|
|
| @classmethod |
| async def _check_node_online(cls, space_host: str) -> bool: |
| try: |
| async with aiohttp.ClientSession( |
| timeout=aiohttp.ClientTimeout(total=5) |
| ) as session: |
| async with session.get(space_host) as response: |
| return response.status == 200 |
| except aiohttp.ClientError: |
| return False |
|
|
|
|
| class ServerState(Node): |
| CACHED: bool |
| TASKS: Optional[Dict[str, Task]] = {} |
| NODES: Optional[list[Node]] |
| DB: str = "https://herokuserver-185316.firebaseio.com/" |
|
|
| def get_master(self) -> Optional[Node]: |
| for node in self.NODES: |
| if node.NODE.MASTER: |
| return node.NODE |
| return None |
|
|
|
|
| TELEGRAM_SESSION = os.environ.get("TELEGRAM_SESSION", "RANDOM_STRING") |
| TELEGRAM_SESSION_PYROGRAM = os.environ.get("TELEGRAM_SESSION_PYROGRAM", " ") |
| MASTER_SERVER = bool(os.environ.get("MASTER", 0)) |
| SPACE_HOST = os.environ.get("SPACE_HOST", "RANDOM_STRING") |
|
|
|
|
| SERVER_STATE = ServerState( |
| CACHED=False, MASTER=MASTER_SERVER, SPACE_HOST=SPACE_HOST, NODES=None |
| ) |
|
|
|
|
| bot: TelegramClient = TelegramClient( |
| StringSession(TELEGRAM_SESSION), |
| api_id=870972, |
| api_hash="ce2efaca02dfcd110941be6025e9ac0d", |
| ) |
|
|
|
|
| |
|
|
