#!/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 from .schema import object_id_from_data ROOT_CA_FILE = __file__.replace("soc_collector_cli.py", "data/collector_root_ca.crt") 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(api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Get database info, currently number of documents. :param api_key: The API key. :param base_url: URL to the API. """ req = requests.get(f"{base_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, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Search for documents in the database. :param data: String with either json or path to a json file. :param api_key: The API key. :param base_url: URL to the API. """ search_data = json_load_data(data) req = requests.post( f"{base_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, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Delete a document in the DB. :param data: id or path to a json file containing "_id". :param api_key: The API key. :param base_url: URL to the API. """ if data and isfile(data): json_data = json_load_data(data) object_id = object_id_from_data(json_data) else: object_id = object_id_from_data(data) if object_id is None: print("ERROR: id is not valid") app_exit(1) req = requests.delete( f"{base_url}/sc/v0/{str(object_id)}", 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: {str(object_id)}") 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, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> 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. :param api_key: The API key. :param base_url: URL to the API. """ json_data = json_load_data(data) object_id = object_id_from_data(json_data) if object_id is None: print("ERROR: Valid '_id' key not in data") app_exit(1) req = requests.put( f"{base_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["_id"]}') def insert_action(data: str, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> 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. :param api_key: The API key. :param base_url: URL to the API. """ json_data = json_load_data(data) if "_id" in json_data: print("ERROR: '_id' key in data") app_exit(1) req = requests.post( f"{base_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["_id"]}') def get_action(data: str, api_key: str, base_url: str = "https://collector-dev.soc.sunet.se:8000") -> None: """Get a document from the database. :param data: id or path to a json file containing "_id". :param api_key: The API key. :param base_url: URL to the API. """ if data and isfile(data): json_data = json_load_data(data) object_id = object_id_from_data(json_data) else: object_id = object_id_from_data(data) if object_id is None: print("ERROR: Invalid id") app_exit(1) req = requests.get( f"{base_url}/sc/v0/{str(object_id)}", 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: id 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: id 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 "COLLECTOR_API_KEY" not in environ: print("Missing 'COLLECTOR_API_KEY' in environment") app_exit(1) api_key = environ["COLLECTOR_API_KEY"] if args.action == "get": get_action(args.data, api_key) elif args.action == "insert": insert_action(args.data, api_key) elif args.action == "replace": replace_action(args.data, api_key) 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, api_key) elif args.action == "search": search_action(args.data, api_key) elif args.action == "info": info_action(api_key) else: print("ERROR: Wrong action") app_exit(1) if __name__ == "__main__": main()