exploitdb.py

#!/usr/bin/env python3

import cmd
import csv
import inspect
import re
import shutil
import socket
import sys
from collections import defaultdict
from dataclasses import dataclass, fields
from pathlib import Path
from urllib.error import URLError
from urllib.request import urlopen
from zipfile import ZipFile

BASE_DIR = Path(__file__).resolve().parent
EXPLOITS_DIR = BASE_DIR / "exploits"
EXPLOITS_DIR_ORIG = BASE_DIR / "exploitdb-main"
EXPLOITS_CSV = EXPLOITS_DIR / "files_exploits.csv"
ARCHIVE_PATH = BASE_DIR / "master.zip"
LATEST_PATH = BASE_DIR / ".latest"
ARCHIVE_URL = (
    "https://gitlab.com/exploit-database/exploitdb/-/archive/main/exploitdb-main.zip"
)

DOWNLOAD_BLOCK_SIZE = 65536  # 64 KiB
NETWORK_TIMEOUT = 30  # seconds


@dataclass(frozen=True, slots=True)
class Exploit:
    id: str
    file: str
    description: str
    date_published: str
    author: str
    type: str
    platform: str
    port: str
    date_added: str
    date_updated: str
    verified: str
    codes: str
    tags: str
    aliases: str
    screenshot_url: str
    application_url: str
    source_url: str


FIELDS: tuple[str, ...] = tuple(f.name for f in fields(Exploit))

_REGEX_TOKEN_RE = re.compile(r'^(\w+):r(["\'])(.*)\2$', re.DOTALL)
_FIELD_TOKEN_RE = re.compile(r"^(\w+):(.*)$", re.DOTALL)


def format_bytes(size: int) -> tuple[float, str]:
    value = float(size)
    for unit in ("B", "KB", "MB", "GB"):
        if value < 1024.0 or unit == "GB":
            return value, unit
        value /= 1024.0
    return value, "GB"  # unreachable


def _strip_outer_quotes(s: str) -> str:
    if len(s) >= 2 and s[0] == s[-1] and s[0] in ('"', "'"):
        return s[1:-1]
    return s


def _tokenize(text: str) -> list[str]:
    """Split on whitespace while keeping quoted sections (and their quotes) intact.

    Quotes are preserved in the output so downstream parsing can still
    recognise regex syntax like `field:r"pattern"`. Unbalanced quotes are
    tolerated: the token runs to the end of the input.
    """
    tokens: list[str] = []
    i, n = 0, len(text)
    while i < n:
        while i < n and text[i].isspace():
            i += 1
        if i >= n:
            break
        start = i
        in_quote: str | None = None
        while i < n:
            c = text[i]
            if in_quote is not None:
                if c == in_quote:
                    in_quote = None
                i += 1
            elif c in ('"', "'"):
                in_quote = c
                i += 1
            elif c.isspace():
                break
            else:
                i += 1
        tokens.append(text[start:i])
    return tokens


def _fatal_network_error(action: str, err: Exception) -> None:
    print(
        f"\n\033[1;31mNetwork error while {action}:\033[0m {err}",
        file=sys.stderr,
    )
    sys.exit(1)


class ExploitSearch(cmd.Cmd):
    prompt: str = "\033[1;31mexploitdb\033[0m\033[1;32m>\033[0m "
    intro: str = (
        "\n \033[40m\033[1;37m-=[ "
        "exploitdb.py - Search exploits from exploit-db.com"
        " ]=-\033[0m\n"
    )
    highlighted_fields_map: dict[str, list[str]] = {
        "id": ["id"],
        "description": ["description"],
        "file": ["type", "platform"],
    }

    def __init__(self, csv_file: Path = EXPLOITS_CSV) -> None:
        self.csv_file: Path = csv_file
        self.exploits: list[Exploit] = []
        self.load_csv()
        self.fields_value_completion: dict[str, set[str]] = {
            "platform": {e.platform for e in self.exploits},
            "type": {e.type for e in self.exploits},
            "port": {e.port for e in self.exploits},
            "verified": {e.verified for e in self.exploits},
        }
        super().__init__()

    @staticmethod
    def get_archive_etag() -> str:
        with urlopen(ARCHIVE_URL, timeout=NETWORK_TIMEOUT) as resp:
            etag = resp.info().get("etag") or ""
        return etag.strip('"')

    @staticmethod
    def download_archive() -> None:
        downloaded_size = 0
        try:
            with (
                urlopen(ARCHIVE_URL, timeout=NETWORK_TIMEOUT) as resp,
                open(ARCHIVE_PATH, "wb") as outfile,
            ):
                while chunk := resp.read(DOWNLOAD_BLOCK_SIZE):
                    outfile.write(chunk)
                    downloaded_size += len(chunk)
                    readable_size, unit = format_bytes(downloaded_size)
                    sys.stdout.write(
                        f"\rDownloading exploits archive... "
                        f"{readable_size:.2f} {unit}   "
                    )
                    sys.stdout.flush()
        except (URLError, socket.timeout) as err:
            _fatal_network_error("downloading exploits archive", err)
        sys.stdout.write("\n")

    def parse_args(self, args: str) -> dict[str, list[tuple[str, str]]]:
        search_args: dict[str, list[tuple[str, str]]] = {"plain": [], "regex": []}
        stripped_args = args.strip()
        if not stripped_args:
            return search_args

        # Shortcut: when there is no `:` anywhere, the whole line is a
        # description phrase search. Strip outer quotes so that both
        # `search wordpress theme` and `search "wordpress theme"` behave
        # the same way.
        if ":" not in stripped_args:
            search_args["plain"].append(
                ("description", _strip_outer_quotes(stripped_args))
            )
            return search_args

        for token in _tokenize(stripped_args):
            regex_match = _REGEX_TOKEN_RE.match(token)
            if regex_match and regex_match.group(1) in FIELDS:
                search_args["regex"].append(
                    (regex_match.group(1), regex_match.group(3))
                )
                continue
            field_match = _FIELD_TOKEN_RE.match(token)
            if field_match:
                field_name = field_match.group(1)
                value = _strip_outer_quotes(field_match.group(2))
                if field_name in FIELDS:
                    search_args["plain"].append((field_name, value))
                    continue
                # Unknown field: reconstruct as a description search so users
                # can still find strings like "foo:bar" inside descriptions.
                search_args["plain"].append(
                    ("description", f"{field_name}:{value}")
                )
                continue
            search_args["plain"].append(
                ("description", _strip_outer_quotes(token))
            )
        return search_args

    def _read_csv_file(self) -> list[Exploit]:
        with open(self.csv_file, newline="") as infile:
            reader = csv.DictReader(infile)
            header = set(reader.fieldnames or ())
            missing = set(FIELDS) - header
            if missing:
                raise RuntimeError(
                    f"CSV is missing expected columns: {sorted(missing)}"
                )
            result: list[Exploit] = []
            for row in reader:
                if not row["port"]:
                    row["port"] = "n/a"
                if not row["platform"]:
                    row["platform"] = "n/a"
                if "//" in row["file"]:
                    row["file"] = row["file"].replace("//", "/")
                result.append(Exploit(**{k: row[k] for k in FIELDS}))
        return result

    def load_csv(self) -> None:
        if not self.csv_file.exists():
            print("Database not found, updating now")
            try:
                latest_etag = self.get_archive_etag()
            except (URLError, socket.timeout) as err:
                _fatal_network_error("checking for database update", err)
                return  # unreachable
            self.updatedb(latest_etag)
        elif not LATEST_PATH.exists():
            print("Version marker missing, attempting update... ", end="", flush=True)
            try:
                latest_etag = self.get_archive_etag()
            except (URLError, socket.timeout) as err:
                print(f"offline ({err}) — using local copy")
            else:
                print("UPDATE FOUND")
                self.updatedb(latest_etag)
        else:
            print("Checking for new database version... ", end="", flush=True)
            try:
                latest_etag = self.get_archive_etag()
            except (URLError, socket.timeout) as err:
                print(f"offline ({err}) — using local copy")
            else:
                current_etag = LATEST_PATH.read_text().strip()
                if latest_etag != current_etag:
                    print("UPDATE FOUND")
                    self.updatedb(latest_etag)
                else:
                    print("OK")

        self.exploits = self._read_csv_file()

    def search(
        self, search_params: str
    ) -> list[tuple[Exploit, dict[str, list[tuple[str, str]]]]]:
        matches: list[tuple[Exploit, dict[str, list[tuple[str, str]]]]] = []
        args = self.parse_args(search_params)
        for exploit in self.exploits:
            if self._exploit_matches(exploit, args):
                matches.append((exploit, args))
        return matches

    @staticmethod
    def _exploit_matches(
        exploit: Exploit, args: dict[str, list[tuple[str, str]]]
    ) -> bool:
        for field_name, pattern in args["plain"]:
            value = getattr(exploit, field_name)
            if pattern.lower() not in value.lower():
                return False
        for field_name, pattern in args["regex"]:
            value = getattr(exploit, field_name)
            if re.search(pattern, value, flags=re.I) is None:
                return False
        return True

    def do_search(self, line: str) -> None:
        """
        search - search the exploits database.

        Type `help search` for the full query syntax and field list.
        """
        for exploit, args in self.search(line):
            flattened_args: dict[str, list[str]] = defaultdict(list)
            for search_type in ("plain", "regex"):
                for k, v in args[search_type]:
                    flattened_args[k].append(v)
            display = {f: getattr(exploit, f) for f in FIELDS}
            regex_fields = {name for name, _ in args["regex"]}
            for field_name, source_fields in self.highlighted_fields_map.items():
                for source in source_fields:
                    for pattern in flattened_args.get(source, []):
                        re_pattern = (
                            pattern if source in regex_fields else re.escape(pattern)
                        )
                        display[field_name] = re.sub(
                            re_pattern,
                            lambda m: f"\033[1;33m{m.group(0)}\033[0m",
                            display[field_name],
                            flags=re.I,
                            count=1,
                        )
            print(f"[{display['id']}] {display['description']} - {display['file']}")
        print("")

    def complete_search(
        self, _text: str, line: str, _begidx: int, _endidx: int
    ) -> list[str]:
        tokens = line.split()
        last_arg = tokens[-1] if tokens else ""
        if not last_arg:
            return [f + ":" for f in FIELDS]
        if ":" not in last_arg:
            return [f + ":" for f in FIELDS if f.startswith(last_arg)]
        field_name, _, pattern = last_arg.partition(":")
        values = self.fields_value_completion.get(field_name)
        if values is None:
            return []
        if pattern:
            return sorted(v for v in values if v.startswith(pattern))
        return sorted(values)

    @staticmethod
    def help_search() -> None:
        field_lines: list[str] = []
        current = "    "
        for i, name in enumerate(FIELDS):
            piece = name + (", " if i < len(FIELDS) - 1 else "")
            if len(current) + len(piece) > 76:
                field_lines.append(current.rstrip())
                current = "    "
            current += piece
        if current.strip():
            field_lines.append(current.rstrip())
        fields_block = "\n".join(field_lines)
        print(
            "\n"
            "search - search the exploits database\n"
            "\n"
            "Usage:\n"
            "    search QUERY\n"
            "\n"
            "Query syntax (multiple terms are combined with AND):\n"
            "    word                      substring match in description\n"
            '    "phrase with spaces"      match the phrase in description\n'
            "    field:value               substring match on the given field\n"
            '    field:"value with space"  same, for values containing spaces\n'
            '    field:r"regex"            regular expression match on the field\n'
            "    field:r'regex'            (single quotes also accepted)\n"
            "\n"
            "All matching is case-insensitive.\n"
            "\n"
            "Available fields:\n"
            f"{fields_block}\n"
            "\n"
            "Examples:\n"
            "    search sudo\n"
            '    search "buffer overflow"\n'
            "    search platform:linux type:remote\n"
            '    search description:"privilege escalation" platform:windows\n'
            '    search codes:r"CVE-2024-\\d+"\n'
            "    search verified:1 type:webapps\n"
        )

    def info(self, exploit_id: str) -> Exploit | None:
        for exploit in self.exploits:
            if exploit.id == exploit_id:
                return exploit
        return None

    def do_info(self, line: str) -> None:
        """
        info - show all metadata for an exploit

        Usage:
            info EXPLOIT_ID

        Displays every non-empty field for the exploit with the given ID,
        grouped into identification, dates, classification, references and
        URLs. Tab-completion is available on exploit IDs.
        """
        result = self.info(line)
        if result is None:
            print(f"No exploit with this ID: {line}\n")
            return

        sections: list[list[tuple[str, str]]] = [
            [
                ("Filename", result.file),
                ("Description", result.description),
            ],
            [
                ("Published", result.date_published),
                ("Added", result.date_added),
                ("Updated", result.date_updated),
            ],
            [
                ("Author", result.author),
                ("Type", result.type),
                ("Platform", result.platform),
                ("Port", result.port),
                ("Verified", "Yes" if result.verified == "1" else "No"),
            ],
            [
                ("CVEs/Codes", result.codes),
                ("Tags", result.tags),
                ("Aliases", result.aliases),
            ],
            [
                ("Screenshot", result.screenshot_url),
                ("Application", result.application_url),
                ("Source", result.source_url),
            ],
        ]
        visible_sections = [
            [(label, value) for label, value in section if value]
            for section in sections
        ]
        visible_sections = [section for section in visible_sections if section]

        label_width = 12
        lines = [
            f" {label:<{label_width}} | {value}"
            for section in visible_sections
            for label, value in section
        ]
        title = f" #{result.id} "
        total_width = max((len(line) for line in lines), default=0)
        total_width = max(total_width, len(title) + 4)

        print(title.center(total_width, "="))
        for i, section in enumerate(visible_sections):
            if i > 0:
                print("")
            for label, value in section:
                print(f" {label:<{label_width}} | {value}")
        print("=" * total_width + "\n")

    def complete_info(
        self, text: str, _line: str, _begidx: int, _endidx: int
    ) -> list[str]:
        if not text:
            return [e.id for e in self.exploits]
        return [e.id for e in self.exploits if e.id.startswith(text)]

    def updatedb(self, etag: str) -> None:
        self.download_archive()
        print("Extracting files...")
        if EXPLOITS_DIR.exists():
            shutil.rmtree(EXPLOITS_DIR)
        if EXPLOITS_DIR_ORIG.exists():
            shutil.rmtree(EXPLOITS_DIR_ORIG)
        with ZipFile(ARCHIVE_PATH) as archive:
            archive.extractall(path=BASE_DIR)
        ARCHIVE_PATH.unlink()
        EXPLOITS_DIR_ORIG.rename(EXPLOITS_DIR)
        EXPLOITS_CSV.chmod(0o644)
        LATEST_PATH.write_text(etag)
        print("OK\n")

    def do_updatedb(self, _line: str) -> None:
        """
        updatedb - download the latest exploits database

        Usage:
            updatedb

        Fetches the current archive from gitlab.com/exploit-database/exploitdb,
        replaces the local exploits/ directory and reloads the in-memory index.
        """
        try:
            etag = self.get_archive_etag()
        except (URLError, socket.timeout) as err:
            _fatal_network_error("checking for database update", err)
            return  # unreachable
        self.updatedb(etag)
        self.exploits = self._read_csv_file()

    def do_show(self, line: str) -> None:
        """
        show - print an exploit's source file

        Usage:
            show EXPLOIT_ID
            show exploits/path/to/exploit.py

        Prints the raw contents of the exploit file. Accepts either the
        numeric exploit ID or the file path as stored in the database.
        Tab-completion supports both forms.
        """
        field = "id" if line.isdigit() else "file"
        match = next(
            (e for e in self.exploits if getattr(e, field) == line), None
        )
        if match is None:
            print(f"Exploit not found: {line}\n")
            return
        sploit_path = EXPLOITS_DIR / match.file
        try:
            print(sploit_path.read_text(errors="replace"))
        except OSError as err:
            print(f"Error reading exploit file: {err}")
        print("")

    def complete_show(
        self, text: str, _line: str, _begidx: int, _endidx: int
    ) -> list[str]:
        completions: list[str] = []
        if not text:
            completions.extend(e.id for e in self.exploits)
            completions.extend(e.file for e in self.exploits)
        else:
            completions.extend(e.id for e in self.exploits if e.id.startswith(text))
            completions.extend(
                e.file for e in self.exploits if e.file.startswith(text)
            )
        return completions

    def do_help(self, arg: str) -> None:
        """List commands, or show detailed help for a given command."""
        if not arg:
            super().do_help(arg)
            return
        custom_help = getattr(self, "help_" + arg, None)
        if custom_help is not None:
            custom_help()
            return
        do_func = getattr(self, "do_" + arg, None)
        doc = getattr(do_func, "__doc__", None)
        if doc:
            print()
            print(inspect.cleandoc(doc))
            print()
            return
        print(self.nohelp % (arg,))

    def do_EOF(self, _line: str) -> bool:
        print()
        return True


def main() -> None:
    es = ExploitSearch()
    while True:
        try:
            es.cmdloop()
            break
        except KeyboardInterrupt:
            es.intro = ""
            print()


if __name__ == "__main__":
    main()