1  from __future__ import annotations
   2  
   3  import csv
   4  import logging
   5  import os
   6  from concurrent.futures import ThreadPoolExecutor, as_completed
   7  from datetime import datetime
   8  from pathlib import Path
   9  from urllib.parse import unquote
  10  
  11  import click
  12  import primp
  13  
  14  from .duckduckgo_search import DDGS
  15  from .utils import _expand_proxy_tb_alias, json_dumps
  16  from .version import __version__
  17  
  18  logger = logging.getLogger(__name__)
  19  
  20  COLORS = {
  21      0: "black",
  22      1: "red",
  23      2: "green",
  24      3: "yellow",
  25      4: "blue",
  26      5: "magenta",
  27      6: "cyan",
  28      7: "bright_black",
  29      8: "bright_red",
  30      9: "bright_green",
  31      10: "bright_yellow",
  32      11: "bright_blue",
  33      12: "bright_magenta",
  34      13: "bright_cyan",
  35      14: "white",
  36      15: "bright_white",
  37  }
  38  
  39  
  40  def _save_data(keywords: str, data: list[dict[str, str]], function_name: str, filename: str | None) -> None:
  41      filename, ext = filename.rsplit(".", 1) if filename and filename.endswith((".csv", ".json")) else (None, filename)
  42      filename = filename if filename else f"{function_name}_{keywords}_{datetime.now():%Y%m%d_%H%M%S}"
  43      if ext == "csv":
  44          _save_csv(f"{filename}.{ext}", data)
  45      elif ext == "json":
  46          _save_json(f"{filename}.{ext}", data)
  47  
  48  
  49  def _save_json(jsonfile: str | Path, data: list[dict[str, str]]) -> None:
  50      with open(jsonfile, "w", encoding="utf-8") as file:
  51          file.write(json_dumps(data))
  52  
  53  
  54  def _save_csv(csvfile: str | Path, data: list[dict[str, str]]) -> None:
  55      with open(csvfile, "w", newline="", encoding="utf-8") as file:
  56          if data:
  57              headers = data[0].keys()
  58              writer = csv.DictWriter(file, fieldnames=headers, quoting=csv.QUOTE_MINIMAL)
  59              writer.writeheader()
  60              writer.writerows(data)
  61  
  62  
  63  def _print_data(data: list[dict[str, str]]) -> None:
  64      if data:
  65          for i, e in enumerate(data, start=1):
  66              click.secho(f"{i}.\t    {'=' * 78}", bg="black", fg="white")
  67              for j, (k, v) in enumerate(e.items(), start=1):
  68                  if v:
  69                      width = 300 if k in ("content", "href", "image", "source", "thumbnail", "url") else 78
  70                      k = "language" if k == "detected_language" else k
  71                      text = click.wrap_text(
  72                          f"{v}", width=width, initial_indent="", subsequent_indent=" " * 12, preserve_paragraphs=True
  73                      )
  74                  else:
  75                      text = v
  76                  click.secho(f"{k:<12}{text}", bg="black", fg=COLORS[j], overline=True)
  77              input()
  78  
  79  
  80  def _sanitize_keywords(keywords: str) -> str:
  81      keywords = (
  82          keywords.replace("filetype", "")
  83          .replace(":", "")
  84          .replace('"', "'")
  85          .replace("site", "")
  86          .replace(" ", "_")
  87          .replace("/", "_")
  88          .replace("\\", "_")
  89          .replace(" ", "")
  90      )
  91      return keywords
  92  
  93  
  94  def _download_file(url: str, dir_path: str, filename: str, proxy: str | None, verify: bool) -> None:
  95      try:
  96          resp = primp.Client(proxy=proxy, impersonate="random", impersonate_os="random", timeout=10, verify=verify).get(
  97              url
  98          )
  99          if resp.status_code == 200:
 100              with open(os.path.join(dir_path, filename[:200]), "wb") as file:
 101                  file.write(resp.content)
 102      except Exception as ex:
 103          logger.debug(f"download_file url={url} {type(ex).__name__} {ex}")
 104  
 105  
 106  def _download_results(
 107      keywords: str,
 108      results: list[dict[str, str]],
 109      function_name: str,
 110      proxy: str | None = None,
 111      threads: int | None = None,
 112      verify: bool = True,
 113      pathname: str | None = None,
 114  ) -> None:
 115      path = pathname if pathname else f"{function_name}_{keywords}_{datetime.now():%Y%m%d_%H%M%S}"
 116      os.makedirs(path, exist_ok=True)
 117  
 118      threads = 10 if threads is None else threads
 119      with ThreadPoolExecutor(max_workers=threads) as executor:
 120          futures = []
 121          for i, res in enumerate(results, start=1):
 122              url = res["image"] if function_name == "images" else res["href"]
 123              filename = unquote(url.split("/")[-1].split("?")[0])
 124              f = executor.submit(_download_file, url, path, f"{i}_{filename}", proxy, verify)
 125              futures.append(f)
 126  
 127          with click.progressbar(  # type: ignore
 128              length=len(futures), label="Downloading", show_percent=True, show_pos=True, width=50
 129          ) as bar:
 130              for future in as_completed(futures):
 131                  future.result()
 132                  bar.update(1)
 133  
 134  
 135  @click.group(chain=True)
 136  def cli() -> None:
 137      """duckduckgo_search CLI tool"""
 138      pass
 139  
 140  
 141  def safe_entry_point() -> None:
 142      try:
 143          cli()
 144      except Exception as ex:
 145          click.echo(f"{type(ex).__name__}: {ex}")
 146  
 147  
 148  @cli.command()
 149  def version() -> str:
 150      print(__version__)
 151      return __version__
 152  
 153  
 154  @cli.command()
 155  @click.option("-k", "--keywords", required=True, help="text search, keywords for query")
 156  @click.option("-r", "--region", default="wt-wt", help="wt-wt, us-en, ru-ru, etc. -region https://duckduckgo.com/params")
 157  @click.option("-s", "--safesearch", default="moderate", type=click.Choice(["on", "moderate", "off"]))
 158  @click.option("-t", "--timelimit", type=click.Choice(["d", "w", "m", "y"]), help="day, week, month, year")
 159  @click.option("-m", "--max_results", type=int, help="maximum number of results")
 160  @click.option("-o", "--output", help="csv, json or filename.csv|json (save the results to a csv or json file)")
 161  @click.option("-d", "--download", is_flag=True, default=False, help="download results. -dd to set custom directory")
 162  @click.option("-dd", "--download-directory", help="Specify custom download directory")
 163  @click.option("-b", "--backend", default="auto", type=click.Choice(["auto", "html", "lite"]))
 164  @click.option("-th", "--threads", default=10, help="download threads, default=10")
 165  @click.option("-p", "--proxy", help="the proxy to send requests, example: socks5://127.0.0.1:9150")
 166  @click.option("-v", "--verify", default=True, help="verify SSL when making the request")
 167  def text(
 168      keywords: str,
 169      region: str,
 170      safesearch: str,
 171      timelimit: str | None,
 172      backend: str,
 173      output: str | None,
 174      download: bool,
 175      download_directory: str | None,
 176      threads: int,
 177      max_results: int | None,
 178      proxy: str | None,
 179      verify: bool,
 180  ) -> None:
 181      """CLI function to perform a text search using DuckDuckGo API."""
 182      data = DDGS(proxy=_expand_proxy_tb_alias(proxy), verify=verify).text(
 183          keywords=keywords,
 184          region=region,
 185          safesearch=safesearch,
 186          timelimit=timelimit,
 187          backend=backend,
 188          max_results=max_results,
 189      )
 190      keywords = _sanitize_keywords(keywords)
 191      if output:
 192          _save_data(keywords, data, "text", filename=output)
 193      if download:
 194          _download_results(
 195              keywords,
 196              data,
 197              function_name="text",
 198              proxy=proxy,
 199              threads=threads,
 200              verify=verify,
 201              pathname=download_directory,
 202          )
 203      if not output and not download:
 204          _print_data(data)
 205  
 206  
 207  @cli.command()
 208  @click.option("-k", "--keywords", required=True, help="keywords for query")
 209  @click.option("-r", "--region", default="wt-wt", help="wt-wt, us-en, ru-ru, etc. -region https://duckduckgo.com/params")
 210  @click.option("-s", "--safesearch", default="moderate", type=click.Choice(["on", "moderate", "off"]))
 211  @click.option("-t", "--timelimit", type=click.Choice(["Day", "Week", "Month", "Year"]))
 212  @click.option("-size", "--size", type=click.Choice(["Small", "Medium", "Large", "Wallpaper"]))
 213  @click.option(
 214      "-c",
 215      "--color",
 216      type=click.Choice(
 217          [
 218              "color",
 219              "Monochrome",
 220              "Red",
 221              "Orange",
 222              "Yellow",
 223              "Green",
 224              "Blue",
 225              "Purple",
 226              "Pink",
 227              "Brown",
 228              "Black",
 229              "Gray",
 230              "Teal",
 231              "White",
 232          ]
 233      ),
 234  )
 235  @click.option("-type", "--type_image", type=click.Choice(["photo", "clipart", "gif", "transparent", "line"]))
 236  @click.option("-l", "--layout", type=click.Choice(["Square", "Tall", "Wide"]))
 237  @click.option(
 238      "-lic",
 239      "--license_image",
 240      type=click.Choice(["any", "Public", "Share", "ShareCommercially", "Modify", "ModifyCommercially"]),
 241  )
 242  @click.option("-m", "--max_results", type=int, help="maximum number of results")
 243  @click.option("-o", "--output", help="csv, json or filename.csv|json (save the results to a csv or json file)")
 244  @click.option("-d", "--download", is_flag=True, default=False, help="download results. -dd to set custom directory")
 245  @click.option("-dd", "--download-directory", help="Specify custom download directory")
 246  @click.option("-th", "--threads", default=10, help="download threads, default=10")
 247  @click.option("-p", "--proxy", help="the proxy to send requests, example: socks5://127.0.0.1:9150")
 248  @click.option("-v", "--verify", default=True, help="verify SSL when making the request")
 249  def images(
 250      keywords: str,
 251      region: str,
 252      safesearch: str,
 253      timelimit: str | None,
 254      size: str | None,
 255      color: str | None,
 256      type_image: str | None,
 257      layout: str | None,
 258      license_image: str | None,
 259      download: bool,
 260      download_directory: str | None,
 261      threads: int,
 262      max_results: int | None,
 263      output: str | None,
 264      proxy: str | None,
 265      verify: bool,
 266  ) -> None:
 267      """CLI function to perform a images search using DuckDuckGo API."""
 268      data = DDGS(proxy=_expand_proxy_tb_alias(proxy), verify=verify).images(
 269          keywords=keywords,
 270          region=region,
 271          safesearch=safesearch,
 272          timelimit=timelimit,
 273          size=size,
 274          color=color,
 275          type_image=type_image,
 276          layout=layout,
 277          license_image=license_image,
 278          max_results=max_results,
 279      )
 280      keywords = _sanitize_keywords(keywords)
 281      if output:
 282          _save_data(keywords, data, function_name="images", filename=output)
 283      if download:
 284          _download_results(
 285              keywords,
 286              data,
 287              function_name="images",
 288              proxy=proxy,
 289              threads=threads,
 290              verify=verify,
 291              pathname=download_directory,
 292          )
 293      if not output and not download:
 294          _print_data(data)
 295  
 296  
 297  @cli.command()
 298  @click.option("-k", "--keywords", required=True, help="keywords for query")
 299  @click.option("-r", "--region", default="wt-wt", help="wt-wt, us-en, ru-ru, etc. -region https://duckduckgo.com/params")
 300  @click.option("-s", "--safesearch", default="moderate", type=click.Choice(["on", "moderate", "off"]))
 301  @click.option("-t", "--timelimit", type=click.Choice(["d", "w", "m"]), help="day, week, month")
 302  @click.option("-res", "--resolution", type=click.Choice(["high", "standart"]))
 303  @click.option("-d", "--duration", type=click.Choice(["short", "medium", "long"]))
 304  @click.option("-lic", "--license_videos", type=click.Choice(["creativeCommon", "youtube"]))
 305  @click.option("-m", "--max_results", type=int, help="maximum number of results")
 306  @click.option("-o", "--output", help="csv, json or filename.csv|json (save the results to a csv or json file)")
 307  @click.option("-p", "--proxy", help="the proxy to send requests, example: socks5://127.0.0.1:9150")
 308  @click.option("-v", "--verify", default=True, help="verify SSL when making the request")
 309  def videos(
 310      keywords: str,
 311      region: str,
 312      safesearch: str,
 313      timelimit: str | None,
 314      resolution: str | None,
 315      duration: str | None,
 316      license_videos: str | None,
 317      max_results: int | None,
 318      output: str | None,
 319      proxy: str | None,
 320      verify: bool,
 321  ) -> None:
 322      """CLI function to perform a videos search using DuckDuckGo API."""
 323      data = DDGS(proxy=_expand_proxy_tb_alias(proxy), verify=verify).videos(
 324          keywords=keywords,
 325          region=region,
 326          safesearch=safesearch,
 327          timelimit=timelimit,
 328          resolution=resolution,
 329          duration=duration,
 330          license_videos=license_videos,
 331          max_results=max_results,
 332      )
 333      keywords = _sanitize_keywords(keywords)
 334      if output:
 335          _save_data(keywords, data, function_name="videos", filename=output)
 336      else:
 337          _print_data(data)
 338  
 339  
 340  @cli.command()
 341  @click.option("-k", "--keywords", required=True, help="keywords for query")
 342  @click.option("-r", "--region", default="wt-wt", help="wt-wt, us-en, ru-ru, etc. -region https://duckduckgo.com/params")
 343  @click.option("-s", "--safesearch", default="moderate", type=click.Choice(["on", "moderate", "off"]))
 344  @click.option("-t", "--timelimit", type=click.Choice(["d", "w", "m", "y"]), help="day, week, month, year")
 345  @click.option("-m", "--max_results", type=int, help="maximum number of results")
 346  @click.option("-o", "--output", help="csv, json or filename.csv|json (save the results to a csv or json file)")
 347  @click.option("-p", "--proxy", help="the proxy to send requests, example: socks5://127.0.0.1:9150")
 348  @click.option("-v", "--verify", default=True, help="verify SSL when making the request")
 349  def news(
 350      keywords: str,
 351      region: str,
 352      safesearch: str,
 353      timelimit: str | None,
 354      max_results: int | None,
 355      output: str | None,
 356      proxy: str | None,
 357      verify: bool,
 358  ) -> None:
 359      """CLI function to perform a news search using DuckDuckGo API."""
 360      data = DDGS(proxy=_expand_proxy_tb_alias(proxy), verify=verify).news(
 361          keywords=keywords, region=region, safesearch=safesearch, timelimit=timelimit, max_results=max_results
 362      )
 363      keywords = _sanitize_keywords(keywords)
 364      if output:
 365          _save_data(keywords, data, function_name="news", filename=output)
 366      else:
 367          _print_data(data)
 368  
 369  
 370  if __name__ == "__main__":
 371      cli(prog_name="ddgs")
 372