183 lines
6.4 KiB
Python
183 lines
6.4 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from userbot.modules.watches import repository
|
|
from userbot.modules.watches.evaluator import (
|
|
Rules,
|
|
Seen,
|
|
build_rules,
|
|
match_chat_scope,
|
|
match_keyword,
|
|
match_new_message,
|
|
)
|
|
from utils.logging import logger
|
|
from utils.read import watches as watches_read
|
|
from utils.read.accounts import self_user_id
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Coroutine
|
|
|
|
import asyncpg
|
|
|
|
ONLINE_QUIET = timedelta(minutes=15)
|
|
PREVIEW_CHARS = 500
|
|
|
|
|
|
def _message_payload(event: str, seen: Seen) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {
|
|
"event": event,
|
|
"chat_id": seen.chat_id,
|
|
"message_id": seen.message_id,
|
|
"sender_id": seen.sender_id,
|
|
"text": (seen.text or "")[:PREVIEW_CHARS] or None,
|
|
}
|
|
if seen.date is not None:
|
|
payload["date"] = seen.date.isoformat()
|
|
return payload
|
|
|
|
|
|
class WatchCache:
|
|
def __init__(self, pool: asyncpg.Pool, account_id: int) -> None:
|
|
self._pool = pool
|
|
self._account_id = account_id
|
|
self.rules = Rules()
|
|
self._self_id: int | None = None
|
|
self._last_seen: dict[int, datetime] = {}
|
|
|
|
async def refresh(self) -> None:
|
|
rows = await watches_read.list_watches(self._pool, self._account_id)
|
|
self.rules = build_rules(rows)
|
|
self._self_id = await self_user_id(self._pool, self._account_id)
|
|
logger.info(f"[green]Watches cached:[/] {len(self.rules)}")
|
|
|
|
def _own(self, seen: Seen) -> bool:
|
|
return seen.outgoing or (
|
|
self._self_id is not None and seen.sender_id == self._self_id
|
|
)
|
|
|
|
async def _guard(self, work: Coroutine[Any, Any, None]) -> None:
|
|
try:
|
|
await work
|
|
except Exception:
|
|
logger.exception("[red]Watch evaluation failed[/]")
|
|
|
|
async def _fire(self, watch_id: int, payload: dict[str, Any], key: str) -> None:
|
|
await watches_read.insert_alert(
|
|
self._pool, self._account_id, watch_id, payload, dedup_key=key
|
|
)
|
|
|
|
async def _keyword(self, event: str, seen: Seen) -> None:
|
|
for rule in self.rules.keyword:
|
|
if match_keyword(rule, seen):
|
|
await self._fire(
|
|
rule.watch_id,
|
|
_message_payload(event, seen),
|
|
f"{seen.chat_id}:{seen.message_id}",
|
|
)
|
|
|
|
async def on_message(self, seen: Seen) -> None:
|
|
if self._own(seen):
|
|
return
|
|
await self._guard(self._on_message(seen))
|
|
|
|
async def _on_message(self, seen: Seen) -> None:
|
|
await self._keyword("message", seen)
|
|
for watch_id, params in self.rules.new_message:
|
|
if match_new_message(params, seen):
|
|
await self._fire(
|
|
watch_id,
|
|
_message_payload("message", seen),
|
|
f"{seen.chat_id}:{seen.message_id}",
|
|
)
|
|
|
|
async def on_edit(self, seen: Seen) -> None:
|
|
if self._own(seen):
|
|
return
|
|
await self._guard(self._on_edit(seen))
|
|
|
|
async def _on_edit(self, seen: Seen) -> None:
|
|
await self._keyword("edit", seen)
|
|
if seen.edit_date is None:
|
|
return
|
|
stamp = int(seen.edit_date.timestamp())
|
|
for watch_id, params in self.rules.edited:
|
|
if match_chat_scope(params, seen):
|
|
await self._fire(
|
|
watch_id,
|
|
_message_payload("edit", seen),
|
|
f"{seen.chat_id}:{seen.message_id}:{stamp}",
|
|
)
|
|
|
|
async def on_transcript(self, chat_id: int, message_id: int, text: str) -> None:
|
|
if self.rules.keyword:
|
|
await self._guard(self._on_transcript(chat_id, message_id, text))
|
|
|
|
async def _on_transcript(self, chat_id: int, message_id: int, text: str) -> None:
|
|
sender_id = await repository.message_sender(
|
|
self._pool, self._account_id, chat_id, message_id
|
|
)
|
|
seen = Seen(chat_id, message_id, sender_id, text)
|
|
if not self._own(seen):
|
|
await self._keyword("transcript", seen)
|
|
|
|
async def on_deleted(self, chat_id: int | None, message_ids: list[int]) -> None:
|
|
if self.rules.deleted and message_ids:
|
|
await self._guard(self._on_deleted(chat_id, message_ids))
|
|
|
|
async def _on_deleted(self, chat_id: int | None, message_ids: list[int]) -> None:
|
|
rows = await repository.deleted_messages(
|
|
self._pool, self._account_id, chat_id, message_ids
|
|
)
|
|
for row in rows:
|
|
seen = Seen(
|
|
row["chat_id"],
|
|
row["message_id"],
|
|
row["sender_id"],
|
|
row["text"],
|
|
date=row["date"],
|
|
)
|
|
if self._own(seen):
|
|
continue
|
|
for watch_id, params in self.rules.deleted:
|
|
if match_chat_scope(params, seen):
|
|
await self._fire(
|
|
watch_id,
|
|
_message_payload("delete", seen),
|
|
f"{seen.chat_id}:{seen.message_id}",
|
|
)
|
|
|
|
async def on_status(
|
|
self, peer_id: int, status: str, at: datetime, last_online: datetime | None
|
|
) -> None:
|
|
watch_ids = [
|
|
watch_id
|
|
for watch_id, params in self.rules.online
|
|
if params.peer_id == peer_id
|
|
]
|
|
if not watch_ids:
|
|
return
|
|
if status != "online":
|
|
self._last_seen[peer_id] = (last_online or at).astimezone(UTC)
|
|
return
|
|
await self._guard(self._on_online(peer_id, at, watch_ids))
|
|
|
|
async def _on_online(
|
|
self, peer_id: int, at: datetime, watch_ids: list[int]
|
|
) -> None:
|
|
stored = await repository.last_seen_before(
|
|
self._pool, self._account_id, peer_id, at
|
|
)
|
|
known = [value for value in (stored, self._last_seen.get(peer_id)) if value]
|
|
previous = max(known, default=None)
|
|
self._last_seen[peer_id] = at
|
|
if previous is not None and at - previous < ONLINE_QUIET:
|
|
return
|
|
session = int((previous or at).timestamp())
|
|
payload: dict[str, Any] = {"event": "online", "peer_id": peer_id}
|
|
if previous is not None:
|
|
payload["since"] = previous.isoformat()
|
|
for watch_id in watch_ids:
|
|
await self._fire(watch_id, payload, f"{peer_id}:{session}")
|