clue-api 1.0.0.dev7__py3-none-any.whl

clue/plugin/helpers/central_server.py

@@ -0,0 +1,27 @@
+# import os
+
+# from clue_client import get_client
+# from clue_client.api.v1 import V1
+# from flask import request
+
+# from clue.common.logging import get_logger
+
+# CENTRAL_SERVER_URL = os.getenv("CENTRAL_API_URL", "http://enrichment-rest.enrichment.svc.cluster.local:5000")
+
+# logger = get_logger(__file__)
+
+
+# def connect_to_central_server(timeout: int | None = 3, retries: int = 3) -> V1:
+#     "Connect to the central server using the clue client"
+#     access_token = request.headers.get("X-Clue-Authorization", None)
+
+#     if access_token:
+#         logger.info("X-Clue-Authorization header specified, using pre-OBO token")
+#     else:
+#         logger.warning("X-Clue-Authorization header not specified, falling back to core Authorization Header")
+#         access_token = request.headers.get("Authorization", " ", type=str).split(" ")[1]
+
+#     if not access_token:
+#         logger.warning("No token specified, continuing with no authentication")
+
+#   return get_client(CENTRAL_SERVER_URL, auth=access_token, version=1, timeout=timeout, retries=retries, logger=logger)

clue/plugin/helpers/email_render.py

@@ -0,0 +1,228 @@
+import base64
+import email
+import email.header
+import io
+import os
+import quopri
+import re
+import tempfile
+import textwrap
+from email.errors import HeaderParseError
+from email.message import Message
+from tempfile import NamedTemporaryFile
+from typing import cast
+
+# TODO: Better handle these specific imports in dependency management
+import imgkit
+from bs4 import BeautifulSoup
+from cart import unpack_stream
+from PIL import Image
+
+from clue.common.exceptions import ClueException, ClueRuntimeError
+from clue.common.logging import get_logger
+from clue.models.results.image import ImageResult
+
+logger = get_logger(__file__)
+
+
+TEXT_TYPES = ["text/plain", "text/html"]
+IMAGE_TYPES = ["image/gif", "image/jpeg", "image/png"]
+
+
+def append_images(images):
+    "Concatenate images together"
+    try:
+        bg_color = (255, 255, 255)
+        widths, heights = zip(*(i.size for i in images))
+
+        new_width = max(widths)
+        new_height = sum(heights)
+        new_im = Image.new("RGB", (new_width, new_height), color=bg_color)
+        offset = 0
+        for im in images:
+            x = 0
+            new_im.paste(im, (x, offset))
+            offset += im.size[1]
+        return new_im
+    except Exception as e:
+        raise ClueException("Error when appending images") from e
+
+
+def get_header_data(msg: Message, header: str) -> str:
+    "Decode the value of a given header"
+    try:
+        decode = email.header.decode_header(msg[header])[0]
+
+        if isinstance(decode[0], bytes):
+            value = decode[0].decode()
+        else:
+            value = str(decode[0])
+    except HeaderParseError:
+        logger.warning("Could not parse header [%s], defaulting to Unknown", header)
+        value = "<Unknown>"
+    logger.info("%s: %s", header, value)
+    return value.replace("<", "&lt;").replace(">", "&gt;")
+
+
+def filter_elements(payload: str) -> str:
+    "Filter external links from html"
+    logger.info("Performing trimming of external fonts/styles")
+    soup = BeautifulSoup(payload, "html.parser")
+
+    logger.debug("Checking for link elements")
+    for link in soup.select("link"):
+        logger.debug("Removing link tag: %s", link["href"])
+        link.decompose()
+
+    logger.debug("Checking for style elements with imports")
+    for style in soup.select("style"):
+        if style.string and "@import" in style.string:
+            _import_match = re.search(r"(@import.+;)", style.string)
+            if _import_match:
+                logger.debug("Removing import: %s", _import_match.group(1))
+            style.string = re.sub(r"@import.+;", "", style.string).strip()
+
+            if not style.string:
+                style.decompose()
+
+    logger.debug("Checking for script tags")
+    for script in soup.select("script"):
+        logger.debug("Removing script")
+        script.decompose()
+
+    return cast(str, soup.prettify())
+
+
+def process_eml(data, output_dir, load_images=False):  # noqa: C901
+    "Process the email (bytes), extract MIME parts and useful headers. Generate a PNG picture of the mail"
+    logger.debug("Beginning eml processing")
+
+    try:
+        msg = email.message_from_bytes(data)
+        date_field = get_header_data(msg, "Date")
+        from_field = get_header_data(msg, "From")
+        to_field = get_header_data(msg, "To")
+        subject_field = get_header_data(msg, "Subject")
+        id_field = get_header_data(msg, "Message-Id")
+
+        imgkit_options = {"load-error-handling": "skip", "no-images": None}
+
+        images_list = []
+
+        # Build a first image with basic mail details
+        headers = textwrap.dedent(f"""
+        <table width="100%%">
+            <tr><td align="right"><b>Date:</b></td><td>{date_field}</td></tr>
+            <tr><td align="right"><b>From:</b></td><td>{from_field}</td></tr>
+            <tr><td align="right"><b>To:</b></td><td>{to_field}</td></tr>
+            <tr><td align="right"><b>Subject:</b></td><td>{subject_field}</td></tr>
+            <tr><td align="right"><b>Message-Id:</b></td><td>{id_field}</td></tr>
+        </table>
+        <hr></p>
+        """)
+
+        try:
+            header_path = NamedTemporaryFile(suffix=".png").name
+            imgkit.from_string(headers, header_path, options=imgkit_options)
+            logger.info("Created headers: %s", header_path)
+            images_list.append(header_path)
+        except Exception:
+            logger.exception("Creation of headers failed.")
+
+        #
+        # Main loop - process the MIME parts
+        #
+        for part in msg.walk():
+            mime_type = part.get_content_type()
+            if part.is_multipart():
+                logger.info("Multipart found, continue")
+                continue
+
+            logger.info("Found MIME part: %s" % mime_type)
+            if mime_type in TEXT_TYPES:
+                try:
+                    # Fix formatting
+                    payload = part.get_payload(decode=True)
+                    payload = re.sub(rb"(\r\n){1,}", b"\r\n", payload)  # type: ignore[arg-type]
+                    payload = payload.replace(b"\r\n", b"<br>")
+                    payload = re.sub(rb"(<br> ){2,}", b"<br><br>", payload)
+
+                    payload = quopri.decodestring(payload).decode("utf-8", errors="ignore")
+                except Exception:
+                    payload = str(quopri.decodestring(part.get_payload(decode=True)))[2:-1]  # type: ignore[arg-type]
+
+                payload = filter_elements(payload)
+
+                try:
+                    payload_path = NamedTemporaryFile(suffix=".png").name
+                    imgkit.from_string(payload, payload_path, options=imgkit_options)
+                    logger.info("Decoded %s" % payload_path)
+                    images_list.append(payload_path)
+                except Exception as e:
+                    logger.warning(f"Decoding this MIME part returned error: {e}")
+
+            elif mime_type in IMAGE_TYPES and load_images:
+                payload = part.get_payload(decode=False)
+                payload_path = NamedTemporaryFile(suffix=".png").name
+                imgdata = base64.b64decode(payload)  # type: ignore[arg-type]
+                try:
+                    with open(payload_path, "wb") as f:
+                        f.write(imgdata)
+                    logger.info("Decoded %s" % payload_path)
+                    images_list.append(payload_path)
+                except Exception as e:
+                    logger.warning(f"Decoding this MIME part returned error: {e}")
+
+        result_image = os.path.join(output_dir, "output.png")
+        if len(images_list) > 0:
+            images = list(map(Image.open, images_list))
+            combo = append_images(images)
+            combo.save(result_image)
+            # Clean up temporary images
+            for i in images_list:
+                os.remove(i)
+            return result_image
+        else:
+            return False
+    except Exception as e:
+        logger.exception("Exception when processing eml")
+        raise ClueException("Error when processing email") from e
+
+
+def render(email_path: str, cart_buffer: io.BytesIO) -> ImageResult | None:
+    "Helper function that, given a buffer containing a carted email, returns an image rendering of it."
+    cart_buffer.seek(0)
+    buf = io.BytesIO()
+    unpack_stream(cart_buffer, buf)
+    buf.seek(0)
+
+    logger.debug("Initializing temporary directory")
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        logger.debug("Temporary directory initialized: %s", tmp_dir)
+
+        process_eml(
+            buf.read(),
+            tmp_dir,
+        )
+
+        error = None
+        if any("output" in s for s in os.listdir(tmp_dir)):
+            previews = [s for s in os.listdir(tmp_dir) if "output" in s]
+            if len(previews) == 0:
+                error = "Target file couldn't be converted to image."
+            elif len(previews) > 1:
+                error = "Target file is generating multiple images."
+        else:
+            error = "Output file does not exist."
+
+        if error:
+            raise ClueRuntimeError(error)
+
+        # There is only 1 rendered image
+        with open(f"{tmp_dir}/{previews[0]}", "rb") as f:
+            return ImageResult(
+                image=f"data:image/png;base64,{base64.b64encode(f.read()).decode("utf-8")}",
+                alt=f"Rendering of {email_path}",
+            )
+
+    logger.warning("No image result generated - returning None")

clue/plugin/helpers/token.py

@@ -0,0 +1,34 @@
+import base64
+import json
+import re
+from typing import Optional, cast
+
+from flask import request
+
+from clue.common.exceptions import InvalidDataException
+
+
+def get_username(token: Optional[str] = None, claims: list[str] | None = None):
+    "Get the username from a given token. Provide an optional list of claims to use"
+    if not token:
+        token = cast(str, request.headers.get("Authorization", None, type=str)).split()[1]
+
+    if not token or "." not in token:
+        raise InvalidDataException("Function requires a JWT to parse username")
+
+    jwt = json.loads(base64.b64decode(token.split(".")[1] + "==").decode())
+
+    if claims is None:
+        claims = ["email", "upn", "unique_name", "preferred_username"]
+
+    username = None
+    for claim in claims:
+        username = jwt.get(claim, None)
+
+        if username:
+            break
+
+    if not username:
+        username = re.sub(r"[^a-z]+", "_", jwt["name"].lower())
+
+    return username

clue/plugin/helpers/trino.py

@@ -0,0 +1,103 @@
+import os
+from typing import Any, Literal, cast
+
+from flask import request
+from trino.auth import JWTAuthentication
+from trino.dbapi import Connection, Cursor, connect
+
+from clue.common.logging import get_logger
+from clue.plugin.helpers.token import get_username
+
+logger = get_logger(__file__)
+
+
+def get_trino_connection(
+    app_name: str,
+    connections: dict[str, Connection],
+    classification: Literal["u"] | Literal["pb"] = "pb",
+    request_timeout: int = 60,
+    max_attempts: int = 2,
+    username_claims: list[str] | None = None,
+    access_token: str | None = None,
+) -> Connection:
+    "Get a trino connection based on the provided JWT"
+    jwt_token: str = access_token or cast(str, request.headers.get("Authorization", None, type=str)).split(" ")[1]
+    if jwt_token not in connections:
+        connections[jwt_token] = connect(
+            http_scheme="https",
+            host=os.environ.get("TRINO_HOST", f"trino.hogwarts.{classification}.azure.chimera.cyber.gc.ca"),
+            port=int(os.environ.get("TRINO_PORT", "443")),
+            user=get_username(jwt_token, claims=username_claims),
+            auth=JWTAuthentication(jwt_token),
+            source=f"clue-{app_name}",
+            max_attempts=max_attempts,
+            request_timeout=request_timeout,
+            # This will stop trino from being bombarded with EXECUTE IMMEDIATE test queries
+            legacy_prepared_statements=False,
+        )
+
+    return connections[jwt_token]
+
+
+def __prepare_query(
+    query: str, where_clause: str, limit: int | None, entries: list[list[str]] | list[str]
+) -> tuple[str, list[str]]:
+    num_where_args = len([character for character in list(where_clause) if character == "?"])
+    if num_where_args == 1 and any(isinstance(entry, list) for entry in entries):
+        logger.error(
+            "Invalid number of arguments provided for where clause. The where clause has one "
+            "?, but you provided a list of arguments."
+        )
+        return "invalid", []
+    elif num_where_args > 1 and not all(isinstance(entry, list) for entry in entries):
+        logger.error(
+            "Invalid number of arguments provided for where clause. The where clause has %s "
+            "?, but you did not provide a list of arguments.",
+            num_where_args,
+        )
+        return "invalid", []
+    elif num_where_args > 1 and not all(len(entry) == num_where_args for entry in entries):
+        logger.error(
+            "Invalid number of arguments provided for where clause. The where clause has %s "
+            "?, but you provided a list of arguments of length %s.",
+            num_where_args,
+            len(entries[0]),
+        )
+        return "invalid", []
+
+    final_query = query.strip()
+    if not final_query.strip().lower().endswith("where"):
+        final_query = final_query.strip() + " WHERE "
+    else:
+        final_query += " "
+
+    values = []
+    for entry in entries:
+        if not isinstance(entry, list):
+            values.append(entry)
+        else:
+            values += entry
+
+        final_query += f"({where_clause}) OR "
+
+    final_query = final_query[:-4]
+
+    if limit is not None:
+        final_query += f" LIMIT {limit}"
+
+    return final_query, values
+
+
+def execute_bulk_query(
+    cur: Cursor, query: str, where_clause: str, limit: int | None, entries: list[list[str]] | list[str]
+) -> list[dict[str, Any]]:
+    "Build a bulk SQL query based on a main query, a template where clause, and a list of entries."
+    final_query, values = __prepare_query(query, where_clause, limit, entries)
+
+    cur.execute(final_query, values)
+
+    results: list[dict[str, Any]] = []
+    for row in cur.fetchall():
+        results.append(dict(zip([desc.name for desc in cur.description], row)))
+
+    return cur.fetchall()

clue/plugin/interactive.py

@@ -0,0 +1,270 @@
+import importlib
+import inspect
+import json
+import os
+import sys
+import textwrap
+from pathlib import Path
+from typing import Callable
+from urllib.parse import quote_plus
+
+from flask.testing import FlaskClient
+from termcolor import colored
+
+from clue.plugin import CluePlugin
+
+TESTABLE_FUNCTIONS = [
+    ("get_actions", None),
+    ("execute_action", "run_action"),
+    ("get_fetchers", None),
+    ("execute_fetcher", "run_fetcher"),
+    ("get_type_names", None),
+    ("lookup", "enrich"),
+    ("bulk_lookup", "enrich"),
+    ("liveness", None),
+    ("readyness", None),
+]
+
+
+def success(*messages: str):
+    "Print success message"
+    print(f"[{colored("success", "green")}]", *messages)
+
+
+def warn(*messages: str):
+    "Print error message"
+    print(f"[{colored("warn", "yellow")}]", *messages)
+
+
+def error(*messages: str):
+    "Print error message"
+    print(f"[{colored("error", "red")}]", *messages)
+
+
+def info(*messages: str):
+    "Print info message"
+    print(f"[{colored("info", "cyan")}]", *messages)
+
+
+class CustomTestClient(FlaskClient):
+    "Custom test client to inject authorization headers"
+
+    def open(self, *args, buffered=False, follow_redirects=False, **kwargs):
+        "Overriden open funciton to inject auth header"
+        headers = kwargs.setdefault("headers", {})
+
+        if "CLUE_ACCESS_TOKEN" in os.environ:
+            info("Clue access token in env, setting Authorization header")
+            headers["Authorization"] = f"Bearer {os.environ["CLUE_ACCESS_TOKEN"]}"
+        else:
+            warn("Missing access token, skipping authorization header.")
+
+        return super().open(*args, buffered=buffered, follow_redirects=follow_redirects, **kwargs)
+
+
+def filter_members(member, current_module):
+    "Get a filtered list of members exported by a given application"
+    member_module = inspect.getmodule(member)
+
+    if member_module is None:
+        return False
+
+    if member_module == current_module:
+        return True
+
+    if not member_module.__name__.startswith("clue"):
+        return False
+
+    return True
+
+
+def test_function(plugin: CluePlugin, fn_id: str, fn: Callable):  # noqa: C901
+    "test a function"
+    info(f"Executing test functionality for {fn_id}")
+
+    plugin.app.test_client_class = CustomTestClient
+
+    for rule in plugin.app.url_map.iter_rules():
+        if rule.endpoint != fn_id:
+            continue
+
+        if "GET" in (rule.methods or {}) and "<" not in rule.rule:
+            info("Simple endpoint detected. Running GET")
+            response = plugin.app.test_client().get(rule.rule)
+            info("Response:", json.dumps(response.json, indent=2) if response.json else response.data.decode())
+        elif "GET" in (rule.methods or {}):
+            kwargs: dict[str, str] = {}
+            info(f"{len(rule.arguments)} arguments are necessary. Supply them now:")
+            for argument in sorted(list(rule.arguments)):
+                kwargs[argument] = quote_plus(quote_plus(input(f"{argument}: ")))
+
+            with plugin.app.test_request_context():
+                path = plugin.app.url_for(fn_id, **kwargs)  # type: ignore[arg-type]
+                info(f"Making request to path {path}")
+
+                response = plugin.app.test_client().get(path)
+
+                if response.status_code > 299:
+                    error(
+                        (response.json or {}).get(
+                            "api_error_message", f"An unknown error occurred. Full response:\n{response.text}"
+                        )
+                    )
+                else:
+                    info("Response:", json.dumps(response.json, indent=2) if response.json else response.data.decode())
+        elif "POST" in (rule.methods or {}):
+            kwargs: dict[str, str] = {}
+            if "<" in rule.rule:
+                info(f"{len(rule.arguments)} arguments are necessary. Supply them now:")
+                for argument in sorted(list(rule.arguments)):
+                    kwargs[argument] = quote_plus(quote_plus(input(f"{argument}: ")))
+
+            info(
+                "Endpoint requires POST data. You can probide a JSON file for this data. "
+                f"Provide a path relative to {os.getcwd()} or an absolute path."
+            )
+            json_path = Path(os.getcwd()) / input("Path to JSON: ").strip()
+
+            if not json_path.exists() or json_path.is_dir():
+                error(f"Provided path {json_path} is invalid or is a directory.")
+                return
+
+            with json_path.open("r") as _data, plugin.app.test_request_context():
+                try:
+                    post_data = json.load(_data)
+                except json.JSONDecodeError:
+                    error(f"The file data in {json_path} is not valid JSON.")
+                    return
+
+                api_path = rule.rule if "<" not in rule.rule else plugin.app.url_for(fn_id, **kwargs)  # type: ignore[arg-type]
+
+                info(f"Submitting POST request to {api_path}:\n{post_data}")
+
+                response = plugin.app.test_client().post(
+                    api_path,
+                    data=json.dumps(post_data),
+                    headers={"Content-Type": "application/json"},
+                )
+
+                if response.status_code > 299:
+                    error(
+                        (response.json or {}).get(
+                            "api_error_message", f"An unknown error occurred. Full response:\n{response.text}"
+                        )
+                    )
+                else:
+                    info("Response:", json.dumps(response.json, indent=2) if response.json else response.data.decode())
+
+
+def main():  # noqa: C901
+    "main interactive loop"
+    os.environ["ENABLE_CACHE"] = "false"
+
+    plugin_name = None
+    if len(sys.argv) > 1:
+        plugin_name = sys.argv[1]
+
+        if not (Path(__file__).parent.parent.parent / "plugins" / plugin_name).exists():
+            error(f"Plugin {plugin_name} does not exist.")
+            plugin_name = None
+
+    while plugin_name is None:
+        plugin_name = input("What plugin do you want to interact with?\n> ")
+        if not (Path(__file__).parent.parent.parent / "plugins" / plugin_name).exists():
+            error(f"Plugin {plugin_name} does not exist.")
+            plugin_name = None
+
+    try:
+        _module = importlib.import_module(f"plugins.{plugin_name}.app")
+        success(f"Initializing plugin {plugin_name} for interactivity")
+    except Exception:
+        error(f"Initializing plugin {plugin_name} for interactivity")
+        raise
+
+    plugin: CluePlugin | None = None
+    for key, member in inspect.getmembers(_module, predicate=lambda _m: filter_members(_m, _module)):
+        if isinstance(member, CluePlugin):
+            success(f"Plugin found exported as member {key}")
+            plugin = member
+            break
+
+    if plugin is None:
+        error("CluePlugin object is not exported from this module!")
+        return
+
+    plugin.cache = None
+
+    functions: list[tuple[str, Callable]] = []
+
+    for attribute in dir(plugin):
+        test_entry = next((entry for entry in TESTABLE_FUNCTIONS if entry[0] == attribute), None)
+        if test_entry is None:
+            continue
+
+        fn = plugin.__getattribute__(attribute)
+        if fn is None:
+            continue
+
+        if test_entry[1] is not None:
+            helper_fn = plugin.__getattribute__(test_entry[1])
+
+            if helper_fn is None:
+                continue
+
+        functions.append((attribute, fn))
+
+    choice: int | None = None
+
+    print(
+        textwrap.dedent("""
+        Clue Plugin Development Script
+
+        This script will help you test various aspects of your plugin interactively.
+        """),
+    )
+
+    if "CLUE_ACCESS_TOKEN" not in os.environ:
+        warn(
+            textwrap.dedent("""
+            Environment variable CLUE_ACCESS_TOKEN not set!
+
+            It is highly likely your plugin will not work if it connects to an external service.
+
+            You'll need to generate a matching token. If on jupyhub, run:
+
+            export CLUE_ACCESS_TOKEN=$(python -c "from hogwarts.auth.vault.credentials import ClueTokenCredential; print(ClueTokenCredential().get_token().token)")
+
+            If you need a different token use the corresponding credential (i.e., trino -> TrinoTokenCredential, howler -> HowlerTokenCredential).
+
+            If you're developing on your local machine, talk to Matthew Rafuse <Matthew.Rafuse@cyber.gc.ca> on Teams.
+            """).strip()  # noqa: E501
+        )
+
+    while choice is None:
+        print("\nAvailable functions:")
+
+        for i in range(len(functions)):
+            print(f"{i + 1}) {' '.join(word.capitalize() for word in functions[i][0].split("_"))}")
+        print(f"{len(functions) + 1}) Quit")
+
+        action = input("\nEnter a selection: ")
+
+        try:
+            choice = int(action)
+
+            if choice > len(functions) + 1:
+                error(f"Invalid choice, choose option between 1 - {len(functions)}.")
+                choice = None
+        except ValueError:
+            error(f"Invalid integer, choose option between 1 - {len(functions)}.")
+
+        if choice is not None and choice <= len(functions):
+            test_function(plugin, *functions[choice - 1])
+            choice = None
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\rExiting!" + " " * 80)

clue/plugin/models.py

@@ -0,0 +1,19 @@
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from clue.models.network import QueryEntry
+
+
+class BulkEntry(BaseModel):
+    "Bulk plugin response for selectors"
+
+    error: str | None = Field(
+        description="Error message returned by plugin",
+        default=None,
+        examples=["An error occurred when enriching the data.", None],
+    )
+    items: list[QueryEntry] = Field(description="List of results from the plugin", default=[])
+    raw_data: Any | None = Field(default=None, description="The raw data for the given results")
+
+    model_config = ConfigDict(validate_assignment=True)