outside-fetcher/__main__.py
2019-11-03 22:15:52 +01:00

94 lines
2.4 KiB
Python

import typing
import pyrogram
import aiohttp.web
import asyncio
import multiprocessing
import config
from async_worker.async_worker import AsyncTaskScheduler
from tasks.channel_history import ChannelHistoryReadTask
class ClientRefCount:
client: pyrogram.Client
count: int = 0
def __init__(self, api_id: int, api_hash: str, name: str):
self.client = pyrogram.Client(name, api_id, api_hash)
async def start(self):
await self.client.start()
def get_client(self) -> pyrogram.Client:
self.count += 1
return self.client
def __gt__(self, other: 'ClientRefCount'):
return self.count < other.count
class ClientRoundRobin:
clients: typing.List[ClientRefCount]
def __init__(self):
self.clients = []
def get_client(self) -> pyrogram.Client:
return min(self.clients).get_client()
async def add_client(self, api_id: int, api_hash: str, name: str):
client = ClientRefCount(api_id, api_hash, name)
self.clients.append(client)
await client.start()
class FetcherAPI:
clients: ClientRoundRobin
scheduler: AsyncTaskScheduler
already_watching: typing.List[str]
def __init__(self):
self.clients = ClientRoundRobin()
self.scheduler = AsyncTaskScheduler()
self.already_watching = []
async def setup(self):
for name in config.sessions:
await self.clients.add_client(config.app_id, config.app_hash, name)
app = aiohttp.web.Application()
app.add_routes([aiohttp.web.get("/tasks/watch/add/{username}", self.add_channel)])
await asyncio.gather(
aiohttp.web._run_app(app, host=config.listen_host, port=config.listen_port),
*(
self.scheduler.loop()
for _ in range(multiprocessing.cpu_count())
)
)
async def add_channel(self, request):
user = request.match_info["username"]
if user in self.already_watching:
return aiohttp.web.Response(status=400)
self.already_watching.append(user)
client = self.clients.get_client()
task = ChannelHistoryReadTask()
peer = await client.resolve_peer(user)
task.setup(client, peer, config.webhook)
await self.scheduler.submit(task)
return aiohttp.web.Response(status=200)
if __name__ == "__main__":
api = FetcherAPI()
_loop = asyncio.get_event_loop()
_loop.run_until_complete(api.setup())