This repository has no description
23 kB
585 lines
1#!/usr/bin/env python3
2"""
3Kagi search script using session token authentication.
4Scrapes search results from Kagi web interface.
5"""
6
7import json
8import subprocess
9import sys
10import os
11from pathlib import Path
12from typing import List, Dict, Optional, Any
13from urllib.request import Request, HTTPCookieProcessor, build_opener
14from urllib.parse import urlencode
15from urllib.error import URLError, HTTPError
16from http.cookiejar import CookieJar
17import gzip
18from bs4 import BeautifulSoup, Tag
19import time
20from dataclasses import dataclass
21import argparse
22import logging
23
24# Module-level logger
25logger = logging.getLogger(__name__)
26
27
28def colorize(text: str, color: str = "", bold: bool = False, dim: bool = False) -> str:
29 """
30 Colorize text if colors are enabled (TTY detected and NO_COLOR not set).
31
32 Args:
33 text: The text to colorize
34 color: Color name (red, green, yellow, blue, magenta, cyan, white)
35 bold: Whether to make text bold
36 dim: Whether to make text dim/faint
37
38 Returns:
39 Colorized text or plain text based on environment
40 """
41 # Check if we should use colors
42 if not sys.stdout.isatty() or os.environ.get("NO_COLOR"):
43 return text
44
45 # Color codes
46 colors = {
47 "red": "\033[91m",
48 "green": "\033[92m",
49 "yellow": "\033[93m",
50 "blue": "\033[94m",
51 "magenta": "\033[95m",
52 "cyan": "\033[96m",
53 "white": "\033[97m",
54 }
55
56 # Start with empty escape sequence
57 escape = ""
58
59 # Add color if specified
60 if color and color.lower() in colors:
61 escape = colors[color.lower()]
62
63 # Add bold if requested
64 if bold:
65 escape = "\033[1m" + escape
66
67 # Add dim if requested
68 if dim:
69 escape = "\033[2m" + escape
70
71 # Return plain text if no formatting requested
72 if not escape:
73 return text
74
75 # Return formatted text
76 return f"{escape}{text}\033[0m"
77
78
79def hyperlink(url: str, text: str = "") -> str:
80 """
81 Create a terminal hyperlink if supported (TTY detected and NO_COLOR not set).
82
83 Args:
84 url: The URL to link to
85 text: The text to display (defaults to URL if not provided)
86
87 Returns:
88 Hyperlinked text or plain text based on environment
89 """
90 # Check if we should use hyperlinks (same conditions as colors)
91 if not sys.stdout.isatty() or os.environ.get("NO_COLOR"):
92 return text or url
93
94 # Use the URL as text if no text provided
95 display_text = text or url
96
97 # OSC 8 hyperlink format: ESC]8;;URL ESC\TEXT ESC]8;; ESC\
98 return f"\033]8;;{url}\033\\{display_text}\033]8;;\033\\"
99
100
101@dataclass
102class SearchResult:
103 """Represents a single search result from Kagi."""
104
105 title: str
106 url: str
107 snippet: str
108
109
110@dataclass
111class QuickAnswer:
112 """Represents a Kagi Quick Answer response."""
113
114 html: str
115 markdown: str
116 raw_text: str
117 references: List[Dict[str, Any]]
118
119
120class KagiSearch:
121 """Kagi search client using session token authentication."""
122
123 def __init__(self, session_token: Optional[str] = None, config_path: Optional[str] = None):
124 """
125 Initialize Kagi search client.
126
127 Args:
128 session_token: Kagi session token (if not provided, will load from config)
129 config_path: Path to config file (defaults to ~/.config/kagi/config.json)
130 """
131 self.base_url = "https://kagi.com"
132 self.user_agent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
133
134 # Load config and get session token
135 if not session_token:
136 config = self._load_config(config_path)
137 session_token = self._get_session_token(config)
138
139 self.session_token = session_token
140
141 # Set up cookie handling
142 self.cookie_jar = CookieJar()
143 self.opener = build_opener(HTTPCookieProcessor(self.cookie_jar))
144
145 # Authenticate with token
146 self._authenticate()
147
148 def _load_config(self, config_path: Optional[str] = None) -> Dict[str, Any]:
149 """Load configuration from file."""
150 if not config_path:
151 config_file = Path.home() / ".config" / "kagi" / "config.json"
152 else:
153 config_file = Path(config_path)
154
155 if not config_file.exists():
156 # Create default config
157 config_file.parent.mkdir(parents=True, exist_ok=True)
158 default_config = {
159 "password_command": "rbw get kagi-session-link",
160 "timeout": 30,
161 "max_retries": 5,
162 }
163 with open(config_file, "w") as f:
164 json.dump(default_config, f, indent=2)
165 msg = colorize(f"Created default config at {config_file}", color="green")
166 print(msg, file=sys.stderr)
167 return default_config
168
169 with open(config_file, "r") as f:
170 data: Dict[str, Any] = json.load(f)
171 return data
172
173 def _get_session_token(self, config: Dict[str, Any]) -> str:
174 """Get session token using password command from config."""
175 password_command = config.get("password_command", "rbw get kagi-session-link")
176
177 try:
178 # Execute the password command
179 result = subprocess.run(
180 password_command.split(), capture_output=True, text=True, check=True
181 )
182 session_link = result.stdout.strip()
183
184 # Extract token from session link
185 if "token=" in session_link:
186 return session_link.split("token=")[1].split("&")[0]
187 else:
188 # Assume the entire output is the token
189 return session_link
190
191 except subprocess.CalledProcessError as e:
192 error_msg = colorize(f"Error executing password command: {e}", color="red", bold=True)
193 print(error_msg, file=sys.stderr)
194 if e.stderr:
195 stderr_msg = colorize(f"stderr: {e.stderr}", color="red")
196 print(stderr_msg, file=sys.stderr)
197 sys.exit(1)
198 except Exception as e:
199 error_msg = colorize(f"Error getting session token: {e}", color="red", bold=True)
200 print(error_msg, file=sys.stderr)
201 sys.exit(1)
202
203 def _authenticate(self) -> None:
204 """Authenticate with Kagi using session token."""
205 token_url = f"{self.base_url}/html/search?token={self.session_token}"
206 request = Request(token_url)
207 request.add_header("User-Agent", self.user_agent)
208
209 try:
210 response = self.opener.open(request, timeout=30)
211 final_url = response.geturl()
212
213 # Check if we were redirected to home page or html search (successful auth)
214 if final_url in [f"{self.base_url}/", self.base_url, f"{self.base_url}/html/search"]:
215 return # Success
216 elif "/signin" in final_url or "/welcome" in final_url:
217 raise Exception(f"Authentication failed - redirected to {final_url}")
218 except Exception as e:
219 raise Exception(f"Failed to authenticate with token: {e}")
220
221 def search(self, query: str, limit: int = 10) -> List[SearchResult]:
222 """
223 Search Kagi and return results.
224
225 Args:
226 query: Search query
227 limit: Maximum number of results to return
228
229 Returns:
230 List of SearchResult objects
231 """
232 # Construct search URL for HTML version (no token needed, we use cookies)
233 params = {"q": query}
234 search_url = f"{self.base_url}/html/search?{urlencode(params)}"
235
236 max_retries = 5
237 retry_delay = 1.0
238
239 for attempt in range(max_retries):
240 try:
241 # Create request with headers
242 request = Request(search_url)
243 request.add_header("User-Agent", self.user_agent)
244 request.add_header(
245 "Accept",
246 "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
247 )
248 request.add_header("Accept-Language", "en-US,en;q=0.5")
249 request.add_header("DNT", "1")
250 request.add_header("Connection", "keep-alive")
251 request.add_header("Upgrade-Insecure-Requests", "1")
252
253 # Make request using opener (with cookies)
254 response = self.opener.open(request, timeout=30)
255
256 # Check if we're redirected to sign in
257 final_url = response.geturl()
258 if "/signin" in final_url or "/welcome" in final_url:
259 raise Exception(f"Authentication failed - redirected to {final_url}")
260
261 # Read response
262 content = response.read()
263
264 # Handle gzip encoding if present
265 if response.headers.get("Content-Encoding") == "gzip":
266 content = gzip.decompress(content)
267
268 html_content = content.decode("utf-8")
269
270 # Parse results
271 soup = BeautifulSoup(html_content, "html.parser")
272
273 # Find results container
274 results_box = soup.find(class_="results-box")
275 if not results_box or isinstance(results_box, str):
276 if attempt < max_retries - 1:
277 time.sleep(retry_delay)
278 continue
279 raise Exception("No results box found on page")
280
281 # Extract search results
282 results = []
283 if not isinstance(results_box, Tag):
284 raise Exception("Results box is not a Tag element")
285 search_results = results_box.find_all(class_="search-result")
286
287 for result in search_results[:limit]:
288 # Extract title
289 title_elem = result.find(class_="__sri-title")
290 title = ""
291 if title_elem:
292 # Get only the text content, not the UI elements
293 # Remove "More results from this site", "Remove results from this site", etc.
294 title_text = title_elem.get_text(separator=" ", strip=True)
295 # Split by common Kagi UI elements and take the first part
296 for separator in [
297 "More results from",
298 "Remove results from",
299 "Open page in",
300 ]:
301 if separator in title_text:
302 title_text = title_text.split(separator)[0]
303 title = title_text.strip()
304
305 # Extract URL
306 url_box = result.find(class_="__sri-url-box")
307 url = ""
308 if url_box:
309 link = url_box.find("a", href=True)
310 if link:
311 href = link.get("href", "")
312 url = str(href) if href else ""
313
314 # Extract snippet
315 desc_elem = result.find(class_="__sri-desc")
316 snippet = desc_elem.get_text(strip=True) if desc_elem else ""
317
318 if title and url:
319 results.append(SearchResult(title=title, url=url, snippet=snippet))
320
321 if results:
322 return results
323 elif attempt < max_retries - 1:
324 time.sleep(retry_delay)
325 continue
326
327 except (URLError, HTTPError) as e:
328 if attempt < max_retries - 1:
329 time.sleep(retry_delay)
330 continue
331 raise Exception(f"Request failed: {e}")
332
333 return []
334
335 def _get_session_cookie(self) -> str:
336 """Extract kagi_session cookie value for X-Kagi-Authorization header."""
337 for cookie in self.cookie_jar:
338 if cookie.name == "kagi_session":
339 return str(cookie.value) if cookie.value else ""
340 return ""
341
342 def get_quick_answer(self, query: str) -> Optional[QuickAnswer]:
343 """
344 Get Kagi Quick Answer for a query.
345
346 Args:
347 query: Search query
348
349 Returns:
350 QuickAnswer object or None if no answer available
351 """
352 # Construct Quick Answer URL - POST request with query in URL
353 params = {"q": query}
354 quick_answer_url = f"{self.base_url}/mother/context?{urlencode(params)}"
355 logger.debug(f"Quick Answer URL: {quick_answer_url}")
356 logger.info("Fetching Quick Answer...")
357
358 try:
359 # Get session cookie for authorization header
360 session_cookie = self._get_session_cookie()
361
362 # Create POST request with headers matching the working curl command
363 request = Request(quick_answer_url, data=b"", method="POST")
364 request.add_header("User-Agent", self.user_agent)
365 request.add_header("Accept", "application/vnd.kagi.stream")
366 request.add_header("Accept-Language", "en-US,en;q=0.5")
367 request.add_header("Accept-Encoding", "gzip, deflate")
368 request.add_header("Referer", f"{self.base_url}/search?{urlencode(params)}")
369 request.add_header("Origin", self.base_url)
370 request.add_header("Connection", "keep-alive")
371 request.add_header("Content-Length", "0")
372 if session_cookie:
373 request.add_header("X-Kagi-Authorization", session_cookie)
374
375 # Make request using opener (with cookies)
376 logger.debug("Making Quick Answer POST request...")
377 response = self.opener.open(request, timeout=30)
378 logger.debug(f"Response status: {response.getcode()}")
379
380 # Check if we're redirected to sign in
381 final_url = response.geturl()
382 logger.debug(f"Final URL: {final_url}")
383 if "/signin" in final_url or "/welcome" in final_url:
384 raise Exception(f"Authentication failed - redirected to {final_url}")
385
386 # Read response
387 content = response.read()
388 logger.debug(f"Response content length: {len(content)} bytes")
389
390 # Handle gzip encoding if present
391 if response.headers.get("Content-Encoding") == "gzip":
392 logger.debug("Decompressing gzip content")
393 content = gzip.decompress(content)
394
395 # Parse streaming response - lines with prefixes like "hi:", "tokens.json:", "new_message.json:"
396 content_str = content.decode("utf-8")
397 lines = content_str.strip().split("\n")
398 logger.debug(f"Response has {len(lines)} lines")
399
400 final_data = None
401
402 for line in lines:
403 line = line.strip()
404 if not line:
405 continue
406 logger.debug(f"Processing line: {line[:100]}...") # Log first 100 chars
407
408 # Look for the new_message.json line which contains the final answer
409 if line.startswith("new_message.json:"):
410 json_str = line[len("new_message.json:") :]
411 logger.debug(f"Parsing new_message JSON: {json_str[:200]}...")
412 try:
413 # Use JSONDecoder to parse only the first JSON object
414 # This handles cases where there's trailing data after the JSON
415 decoder = json.JSONDecoder()
416 final_data, _ = decoder.raw_decode(json_str)
417 logger.debug("Successfully parsed new_message data")
418 except json.JSONDecodeError as e:
419 logger.error(f"Failed to parse new_message JSON: {e}")
420 logger.debug(f"Full line: {line}")
421
422 if not final_data:
423 logger.debug("No new_message data found in response")
424 return None
425
426 # New response format has fields: md, reply, references_md, etc.
427 markdown = final_data.get("md", "")
428 html = final_data.get("reply", "")
429 references_md = final_data.get("references_md", "")
430
431 # Parse references from references_md (format: [^1]: [Title](URL) (percent%))
432 references = []
433 if references_md:
434 import re
435
436 # Match pattern like: [^1]: [Title](URL) (22%)
437 # Use non-greedy match for URL with lookahead to handle URLs containing parentheses
438 ref_pattern = r"\[\^\d+\]:\s*\[([^\]]+)\]\((.+?)\)\s*\((\d+)%\)"
439 for match in re.finditer(ref_pattern, references_md):
440 references.append(
441 {
442 "title": match.group(1),
443 "url": match.group(2),
444 "contribution": f"{match.group(3)}%",
445 }
446 )
447
448 # If no content, return None
449 if not html and not markdown:
450 logger.debug("No content found in Quick Answer")
451 return None
452
453 logger.debug(f"Quick Answer found with {len(references)} references")
454
455 return QuickAnswer(
456 html=html, markdown=markdown, raw_text=markdown, references=references
457 )
458
459 except Exception as e:
460 # Quick Answer might not be available for all queries
461 logger.debug(f"Quick Answer error: {type(e).__name__}: {e}")
462 return None
463
464
465def main() -> None:
466 """Main entry point for command line usage."""
467 parser = argparse.ArgumentParser(description="Search Kagi using session token")
468 parser.add_argument("query", nargs="?", help="Search query")
469 parser.add_argument(
470 "-n", "--num-results", type=int, default=3, help="Number of link results (default: 3)"
471 )
472 parser.add_argument(
473 "-l", "--links", action="store_true", help="Show search result links (hidden by default)"
474 )
475 parser.add_argument("-t", "--token", help="Session token (overrides config)")
476 parser.add_argument("-c", "--config", help="Config file path")
477 parser.add_argument("-j", "--json", action="store_true", help="Output as JSON")
478 parser.add_argument("-d", "--debug", action="store_true", help="Enable debug logging to stderr")
479
480 args = parser.parse_args()
481
482 # Configure logging - always to stderr
483 log_level = logging.DEBUG if args.debug else logging.WARNING
484 logging.basicConfig(
485 level=log_level,
486 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
487 stream=sys.stderr,
488 )
489
490 # If no query provided, read from stdin
491 if not args.query:
492 if sys.stdin.isatty():
493 parser.error("No query provided")
494 args.query = sys.stdin.read().strip()
495
496 # Initialize client
497 try:
498 client = KagiSearch(session_token=args.token, config_path=args.config)
499 except Exception as e:
500 error_msg = colorize(f"Error initializing Kagi client: {e}", color="red", bold=True)
501 print(error_msg, file=sys.stderr)
502 sys.exit(1)
503
504 # Perform search
505 try:
506 results = client.search(args.query, limit=args.num_results) if args.links else []
507 quick_answer = client.get_quick_answer(args.query)
508 except Exception as e:
509 error_msg = colorize(f"Search failed: {e}", color="red", bold=True)
510 print(error_msg, file=sys.stderr)
511 sys.exit(1)
512
513 # Output results
514 if args.json:
515 output: Dict[str, Any] = {
516 "results": [{"title": r.title, "url": r.url, "snippet": r.snippet} for r in results]
517 }
518 if quick_answer:
519 output["quick_answer"] = {
520 "markdown": quick_answer.markdown,
521 "raw_text": quick_answer.raw_text,
522 "references": quick_answer.references,
523 }
524 print(json.dumps(output, indent=2))
525 else:
526 # Display Quick Answer if available
527 if quick_answer:
528 qa_title = colorize("Quick Answer", color="cyan", bold=True)
529 print(f"\n{qa_title}")
530 print(colorize("─" * 80, color="cyan", dim=True))
531
532 # Display raw text for terminal (cleaner than HTML)
533 if quick_answer.raw_text:
534 print(quick_answer.raw_text)
535 elif quick_answer.markdown:
536 # Fallback to markdown if no raw text
537 print(quick_answer.markdown)
538
539 # Display references
540 if quick_answer.references:
541 print()
542 refs_title = colorize("References:", color="cyan", dim=True)
543 print(refs_title)
544 for i, ref in enumerate(quick_answer.references[:5], 1): # Limit to 5 refs
545 ref_num = colorize(f"[{i}]", color="cyan", dim=True)
546 ref_title = ref.get("title", "")
547 ref_url = ref.get("url", "")
548 if ref_title and ref_url:
549 ref_link = hyperlink(ref_url, colorize(ref_title, color="blue"))
550 print(f" {ref_num} {ref_link}")
551 url_text = colorize(ref_url, color="green", dim=True)
552 url_with_link = hyperlink(ref_url, url_text)
553 print(f" {url_with_link}")
554
555 print(colorize("─" * 80, color="cyan", dim=True))
556
557 # Display search result links only when --links is passed
558 if results:
559 if quick_answer:
560 print() # Extra newline after Quick Answer
561 for i, result in enumerate(results, 1):
562 # Result number in yellow/bold
563 number = colorize(f"{i}.", color="yellow", bold=True)
564 # Title in blue/bold with hyperlink
565 title = colorize(result.title, color="blue", bold=True)
566 title_with_link = hyperlink(result.url, title)
567 print(f"\n{number} {title_with_link}")
568
569 # URL in green with hyperlink
570 url_text = colorize(result.url, color="green")
571 url_with_link = hyperlink(result.url, url_text)
572 print(f" {url_with_link}")
573
574 # Snippet in default color but dim
575 if result.snippet:
576 snippet = colorize(result.snippet, dim=True)
577 print(f" {snippet}")
578
579 if not quick_answer and not results and not args.json:
580 error_msg = colorize("No results found", color="red")
581 print(error_msg, file=sys.stderr)
582
583
584if __name__ == "__main__":
585 main()