playground/plab_search.py

from __future__ import annotations

import argparse
import dataclasses
import datetime
import json
import typing
from pathlib import Path
from pprint import pprint
from urllib.parse import urljoin

import bs4
import httpx


@dataclasses.dataclass()
class SearchResult:
    posted_at: datetime.datetime
    title: str
    size_mb: float
    url: str


class Auth(typing.Protocol):
    def authenticate(self) -> dict[str, str]:
        ...


@dataclasses.dataclass()
class CookieFileAuth:
    file_path: Path

    def authenticate(self) -> dict[str, str]:
        with self.file_path.open() as f:
            return json.load(f)


@dataclasses.dataclass()
class CredentialAuth:
    username: str
    password: str

    def authenticate(self) -> dict[str, str]:
        client = httpx.Client(base_url="https://pornolab.net/")
        res = client.post(
            "/forum/login.php",
            data={
                "login_username": self.username,
                "login_password": self.password,
                "login": "Вход",
            },
            follow_redirects=True,
        )
        res.raise_for_status()
        soup = bs4.BeautifulSoup(res.text, "html.parser")
        username_el = soup.select_one(".topmenu .med")
        if self.username not in username_el.text.strip():
            raise ValueError("failed to login")
        return dict(client.cookies)


@dataclasses.dataclass()
class PersistedCookieAuth:
    cookie_path: Path
    auth: Auth

    def authenticate(self) -> dict[str, str]:
        cookie_auth = CookieFileAuth(self.cookie_path)
        try:
            return cookie_auth.authenticate()
        except (ValueError, FileNotFoundError):
            pass

        cookies = self.auth.authenticate()
        with self.cookie_path.open("w") as f:
            json.dump(cookies, f)
        return cookies


class PlabAuthMiddleware(httpx.Auth):
    def auth_flow(self, request: httpx.Request) -> typing.Generator[httpx.Request, httpx.Response, None]:
        # if no cookies in req
        # fill cookie
        # check response
        # reauth
        # fill cookie

        res = yield request
        return super().auth_flow(request)


class Plab:
    def __init__(self, auth: Auth):
        self.auth = auth
        self.client = httpx.Client(base_url="https://pornolab.net/")

    def search(self, term: str) -> list[SearchResult]:
        self._ensure_session()
        res = self.client.post(
            "/forum/tracker.php",
            data={
                "max": "1",
                "to": "1",
                "nm": f'{term}',
            },
            follow_redirects=True,
        )
        res.raise_for_status()

        if 'Введите ваше имя и пароль' in res.text:
            raise

        soup = bs4.BeautifulSoup(res.text, "html.parser")
        rows = soup.select("#tor-tbl tbody tr")
        items = []
        for row in rows:
            title_link = row.select_one(f'[href*="viewtopic.php?t="]')
            if not title_link:
                continue
            size_mb = (
                int(row.select_one('[href*="dl.php?t="]').find_previous_sibling(name='u').text.strip()) / 1_048_576
            )
            date = datetime.datetime.fromtimestamp(int(row.select_one('[title="Добавлен"] u').text.strip()))
            items.append(
                SearchResult(
                    url=urljoin("https://pornolab.net/forum/", title_link["href"]),
                    title=title_link.text.strip(),
                    size_mb=size_mb,
                    posted_at=date,
                )
            )
        return items

    def _ensure_session(self):
        if self.client.cookies:
            return
        self.client.cookies = self.auth.authenticate()


def dump_for_alfred(items: list[SearchResult]):
    alfreds = []
    for it in items:
        alfreds.append(
            {
                "title": it.title,
                "subtitle": f"{it.size_mb:.1f}MB -- {it.posted_at.isoformat()}",
                "arg": it.url,
                "text": {
                    "copy": it.url,
                    "largetype": f"{it.title}\nPosted at: {it.posted_at.isoformat()}\nSize: {it.size_mb:.1f}MB",
                },
                "match": f"{it.title} {it.posted_at.isoformat()}",
            }
        )
    print(json.dumps({"items": alfreds}, indent=2))


def parse_args():
    arger = argparse.ArgumentParser()
    arger.add_argument('term', help='Term to search for')
    arger.add_argument('--alfred', action='store_true', help='Dump for Alfred')
    return arger.parse_args()


def main():
    args = parse_args()

    p = Plab(
        auth=PersistedCookieAuth(
            cookie_path=Path(__file__).parent / "plab.json",
            auth=CredentialAuth(
                username="zzzp",
                password="3c45HzLMgSB7M5hJ",
            ),
        )
    )

    results = p.search(args.term)

    if args.alfred:
        dump_for_alfred(results)
    else:
        for it in results:
            pprint(it)


if __name__ == "__main__":
    main()