summaryrefslogtreecommitdiff
path: root/skills/web-search/web_search.py
blob: 8124dc9da291dc2dc56459aa5746cbe984c9cc21 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#!/usr/bin/env python3
"""
Search the web via DuckDuckGo's html endpoint and print plain-text results.

Exits 2 on rate-limit/bot-block (retry later), 1 on other errors.
"""

import argparse
import sys
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from html.parser import HTMLParser

ENDPOINT = "https://html.duckduckgo.com/html/"
USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:155.0) Gecko/20100101 Firefox/155.0"


@dataclass
class SearchResult:
    title: str
    url: str
    snippet: str = ""


class ResultsParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.results: list[SearchResult] = []
        # Title text accumulated while inside a result anchor.
        self._title_parts: list[str] | None = None
        self._href: str = ""
        # Index of the result whose snippet we are currently inside.
        self._snippet_for: int | None = None
        self._snippet_parts: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        if tag != "a":
            return
        classes = dict(attrs).get("class", "").split()
        if "result__a" in classes:
            self._title_parts = []
            self._href = dict(attrs).get("href", "")
        elif "result__snippet" in classes:
            self._snippet_for = len(self.results) - 1
            self._snippet_parts = []

    def handle_data(self, data: str) -> None:
        if self._title_parts is not None:
            self._title_parts.append(data)
        elif self._snippet_for is not None:
            self._snippet_parts.append(data)

    def handle_endtag(self, tag: str) -> None:
        if tag != "a":
            return
        if self._title_parts is not None:
            title = " ".join("".join(self._title_parts).split())
            url = self._unwrap_redirect(self._href)
            if title and url:
                self.results.append(SearchResult(title, url))
            self._title_parts = None
        elif self._snippet_for is not None:
            if 0 <= self._snippet_for < len(self.results):
                self.results[self._snippet_for].snippet = " ".join(
                    "".join(self._snippet_parts).split()
                )
            self._snippet_for = None

    @staticmethod
    def _unwrap_redirect(href: str) -> str:
        parsed = urllib.parse.urlparse(href)
        if parsed.netloc.endswith("duckduckgo.com") and parsed.path == "/l/":
            target = urllib.parse.parse_qs(parsed.query).get("uddg", [None])[0]
            return target if target else href
        return href


def fetch(query: str) -> str:
    data = urllib.parse.urlencode({"q": query}).encode()
    request = urllib.request.Request(
        ENDPOINT, data=data, headers={"User-Agent": USER_AGENT}
    )
    with urllib.request.urlopen(request, timeout=15) as response:
        return response.read().decode("utf-8", "replace")


def fail(message: str, exit_code: int = 1) -> None:
    print(f"web_search: {message}", file=sys.stderr)
    sys.exit(exit_code)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    _ = parser.add_argument(
        "-n", type=int, default=5, metavar="N", help="max results (default 5)"
    )
    _ = parser.add_argument("query", nargs="+", help="search terms")
    args = parser.parse_args()
    query = " ".join(args.query)

    try:
        body = fetch(query)
    except urllib.error.HTTPError as err:
        if err.code in (403, 429):
            fail(f"rate-limited or blocked (HTTP {err.code}); try later", exit_code=2)
        fail(f"HTTP {err.code}")
    except OSError as err:
        fail(f"network error: {err}")

    extractor = ResultsParser()
    extractor.feed(body)

    if not extractor.results:
        if "anomaly" in body or "challenge" in body:
            fail("bot challenge served instead of results; try later", exit_code=2)
        print(f"no results for: {query}")
        return

    for number, result in enumerate(extractor.results[: args.n], 1):
        print(f"{number}. {result.title}\n   {result.url}")
        if result.snippet:
            print(f"   {result.snippet}")


if __name__ == "__main__":
    main()