path: root/src/soc_collector/soc_collector_cli.py

#!/usr/bin/env python3
"""Our database module"""
from os.path import isfile
from os import environ
from typing import Dict, Any
from argparse import ArgumentParser, RawTextHelpFormatter
from sys import exit as app_exit
import json
import requests

if "COLLECTOR_API_KEY" not in environ:
    print("Missing 'COLLECTOR_API_KEY' in environment")
    app_exit(1)
API_KEY = environ["COLLECTOR_API_KEY"]

API_URL = "https://collector-dev.soc.sunet.se:8000"
ROOT_CA_FILE = __file__.replace("soc_collector_cli.py", "data/collector_root_ca.crt")


def valid_key(key: str) -> None:
    """Ensure the document key is valid. exit(1) otherwise.

    :param key: The key.
    """
    valid_chars = ["a", "b", "c", "d", "e", "f", "1", "2", "3", "4", "5", "6", "7", "8", "9", "0"]
    if len(key) != 24:
        print(f"ERROR: Invalid key '{key}'")
        app_exit(1)

    for char in key:
        if char not in valid_chars:
            print(f"ERROR: Invalid key '{key}'")
            app_exit(1)


def json_load_data(data: str) -> Dict[str, Any]:
    """Load json from argument, json data or path to json file

    :param data: String with either json or path to a json file.
    :return: json dict as a Dict[str, Any].
    """

    ret: Dict[str, Any]
    try:
        if data and isfile(data):
            with open(data, encoding="utf-8") as f_data:
                ret = json.loads(f_data.read())
        else:
            ret = json.loads(data)
        return ret

    except Exception:  # pylint: disable=broad-except
        print(f"No such file or invalid json data: {data}")
        app_exit(1)


def info_action() -> None:
    """Get database info, currently number of documents."""

    req = requests.get(f"{API_URL}/info", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)

    # Ensure ok status
    req.raise_for_status()

    # Print data
    json_data = json.loads(req.text)
    print(f"Estimated document count: {json_data['Estimated document count']}")


def search_action(data: str) -> None:
    """Search for documents in the database.

    :param data: String with either json or path to a json file.
    """

    search_data = json_load_data(data)

    req = requests.post(
        f"{API_URL}/sc/v0/search", headers={"API-KEY": API_KEY}, json=search_data, timeout=5, verify=ROOT_CA_FILE
    )

    # Ensure ok status
    req.raise_for_status()

    # Print data
    json_data = json.loads(req.text)
    print(json.dumps(json_data["docs"], indent=4))


def delete_action(data: str) -> None:
    """Delete a document in the DB.

    :param data: key or path to a json file containing "_id".
    """
    if data and isfile(data):
        json_data = json_load_data(data)

        if "_id" not in json_data or not isinstance(json_data["_id"], str):
            print("ERROR: Valid '_id' key not in data")
            app_exit(1)
        key: str = json_data["_id"]
    else:
        key = data

    valid_key(key)

    req = requests.delete(f"{API_URL}/sc/v0/{key}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)

    # Check status
    if req.status_code == 404:
        print("ERROR: Document not found")
        app_exit(1)

    # Ensure ok status
    req.raise_for_status()

    print(f"Deleted data OK - key: {key}")


def update_local_action(data: str, update_data: str) -> None:
    """Update keys and their data in this document. Does not modify the database.

    :param data: json blob or path to json file.
    :param update_data: json blob or path to json file.
    """

    json_data = json_load_data(data)
    json_update_data = json_load_data(update_data)

    json_data.update(json_update_data)

    # Print data
    print(json.dumps(json_data, indent=4))


def replace_action(data: str) -> None:
    """Replace the entire document in the database with this document, "_id" must exist as a key.

    :param data: json blob or path to json file, "_id" key must exist.
    """

    json_data = json_load_data(data)

    if "_id" not in json_data or not isinstance(json_data["_id"], str):
        print("ERROR: Valid '_id' key not in data")
        app_exit(1)

    valid_key(json_data["_id"])

    req = requests.put(f"{API_URL}/sc/v0", json=json_data, headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)

    # Check status
    if req.status_code == 404:
        print("ERROR: Document not found")
        app_exit(1)

    # Ensure ok status
    req.raise_for_status()

    json_data = json.loads(req.text)
    print(f'Replaced data OK - key: {json_data["key"]}')


def insert_action(data: str) -> None:
    """Insert a new document into the database, "_id" must not exist in the document.

    :param data: json blob or path to json file, "_id" key must not exist.
    """

    json_data = json_load_data(data)

    if "_id" in json_data:
        print("ERROR: '_id' key in data")
        app_exit(1)

    req = requests.post(
        f"{API_URL}/sc/v0", json=json_data, headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE
    )

    # Ensure ok status
    req.raise_for_status()

    json_data = json.loads(req.text)
    print(f'Inserted data OK - key: {json_data["key"]}')


def get_action(data: str) -> None:
    """Get a document from the database.

    :param data: key or path to a json file containing "_id".
    """
    if data and isfile(data):
        json_data = json_load_data(data)

        if "_id" not in json_data or not isinstance(json_data["_id"], str):
            print("ERROR: Valid '_id' key not in data")
            app_exit(1)
        key: str = json_data["_id"]
    else:
        key = data

    valid_key(key)

    req = requests.get(f"{API_URL}/sc/v0/{key}", headers={"API-KEY": API_KEY}, timeout=5, verify=ROOT_CA_FILE)

    # Check status
    if req.status_code == 404:
        print("ERROR: Document not found")
        app_exit(1)

    # Ensure ok status
    req.raise_for_status()

    # Print data
    json_data = json.loads(req.text)
    print(json.dumps(json_data["doc"], indent=4))


def main() -> None:
    """Main function."""
    parser = ArgumentParser(formatter_class=RawTextHelpFormatter, description="SOC Collector CLI")
    parser.add_argument(
        "action",
        choices=["info", "search", "get", "insert", "replace", "update_local", "delete"],
        help="""Action to take

    info: Show info about the database, currently number of documents.

    search: json blog OR path to file. skip defaults to 0 and limit to 25.
    '{"filter": {"port": 111, "ip": "192.0.2.28"}, "skip": 0, "limit": 25}' OR ./search_data.json
    '{"filter": {"asn_country_code": "SE", "result": {"$exists": "cve_2015_0002"}}}'
    '{"filter": {}}'

    get: key OR path to document using its "_id".
    637162378c92893fff92bf7e OR ./data.json

    insert: json blob OR path to file. Document MUST NOT contain "_id".
    '{json_blob_here...}' OR ./data.json

    replace: json blob OR path to file. Document MUST contain "_id".
    '{json_blob_here...}' OR ./updated_document.json

    update_local: json blob OR path to file json blob or path to file with json to update with.
    This does NOT send data to the database, use replace for that.
    1st ARG: '{json_blob_here...}' OR ./data.json 2th ARG:'{"port": 555, "some_key": "some_data"}' OR ./data.json

    delete: key OR path to file using its "_id".
    637162378c92893fff92bf7e OR ./data.json


    Pro tip, in ~/.bashrc:

    _soc_collector_cli()
    {
        local cur prev words cword
        _init_completion || return

        local commands command

        commands='info search get insert replace update_local delete'

        if ((cword == 1)); then
            COMPREPLY=($(compgen -W "$commands" -- "$cur"))
        else

        case $prev in
            search)
                COMPREPLY="'{\\"filter\\": {\\"asn_country_code\\": \\"SE\\", \\"ip\\": \\"192.0.2.28\\"}}'"
                return
                ;;
        esac
        command=${words[1]}
        _filedir
        return
        fi
    }
    complete -F _soc_collector_cli -o default soc_collector_cli

    """,
    )
    parser.add_argument("data", nargs="?", help="json blob or path to file")
    parser.add_argument(
        "extra_data",
        nargs="?",
        default=None,
        help="json blob or path to file, only used for local_edit",
    )

    args = parser.parse_args()

    if args.action == "get":
        get_action(args.data)
    elif args.action == "insert":
        insert_action(args.data)
    elif args.action == "replace":
        replace_action(args.data)
    elif args.action == "update_local" and args.extra_data is not None:
        update_local_action(args.data, args.extra_data)
    elif args.action == "delete":
        delete_action(args.data)
    elif args.action == "search":
        search_action(args.data)
    elif args.action == "info":
        info_action()
    else:
        print("ERROR: Wrong action")
        app_exit(1)


if __name__ == "__main__":
    main()