import demistomock as demisto from CommonServerPython import * from CommonServerUserPython import * """ IMPORTS """ from datetime import datetime from pymongo import MongoClient """ GLOBALS/PARAMS """ # Get Credentials USERNAME = demisto.params().get("credentials").get("identifier") PASSWORD = demisto.params().get("credentials").get("password") # Get Server URI = demisto.params().get("uri") # Get Database DATABASE = demisto.params().get("database") USE_SSL = demisto.params().get("use_ssl", False) INSECURE = demisto.params().get("insecure", False) TIMEOUT = 5000 if INSECURE and not USE_SSL: raise DemistoException('"Trust any certificate (not secure)" must be ticked with "Use TLS/SSL secured connection"') if not INSECURE and not USE_SSL: # Connect to MongoDB - Need to add credentials and lock down MongoDB (add auth) CLIENT = MongoClient( # type: ignore[var-annotated] URI, username=USERNAME, password=PASSWORD, authSource=DATABASE, authMechanism="SCRAM-SHA-1", ssl=USE_SSL, socketTimeoutMS=TIMEOUT, ) else: CLIENT = MongoClient( URI, username=USERNAME, password=PASSWORD, authSource=DATABASE, authMechanism="SCRAM-SHA-1", ssl=USE_SSL, tlsAllowInvalidCertificates=INSECURE, socketTimeoutMS=TIMEOUT, ) DB = CLIENT[DATABASE] # Set Collection COLLECTION_NAME = demisto.params().get("collection") COLLECTION = DB[COLLECTION_NAME] def test_module(): """Check DB Status""" if CLIENT.server_info().get("ok") == 1.0: return "ok", {}, {} return "MongoDB Server Error", {}, {} def write_log_json(): """Gather Args, form json document, write document to MondoDB""" investigation = demisto.investigation() investigation_id = investigation.get("id") investigation_user = investigation.get("user") timestamp = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S+00:00") id_ = demisto.args().get("id", investigation_id) playbook = demisto.args().get("playbook") action = demisto.args().get("action") user = demisto.args().get("user", investigation_user) message = demisto.args().get("message") logjson = {"timestamp": timestamp, "id": id_, "playbook": playbook, "action": action, "user": user, "message": message} # Add json to the document collection in MondoDB result = COLLECTION.insert_one(logjson) entry_id = result.inserted_id context = { "EntryID": str(entry_id), "Timestamp": timestamp, "ID": id_, "Playbook": playbook, "Action": action, "User": user, "Message": message, } ec = {"MongoDB.Entry(val.EntryID === obj.EntryID)": context} return "MongoDB Log - 1 document/record added", ec, {} def read_log_json(): """Get all log documents/records from MondoDB""" limit = int(demisto.args().get("limit")) # Point to all the documents doc_count = COLLECTION.count_documents(filter={}, limit=limit) cursor = COLLECTION.find({}, {"_id": False}) # Create an empty log list entries = [] # Iterate through those documents if doc_count > 0: for i in cursor: # Append log entry to list entries.append(i) return_json = {COLLECTION_NAME: entries} human_readable = tableToMarkdown( f'The log documents/records for collection "{COLLECTION_NAME}"', return_json.get(COLLECTION_NAME) ) return human_readable, {}, {} return "MongoDB - no documents/records - Log collection is empty", {}, {} def num_log_json(): """Get a count of all log documents/records from MondoDB""" # Point to the documents doc_count = COLLECTION.count_documents(filter={}) human_readable = f"The count of log documents/records is {doc_count!s}" return human_readable, {}, {} def main(): LOG(f"Command being called is {demisto.command()}") try: if demisto.command() == "test-module": # This is the call made when pressing the integration test button. test_module() return_outputs(*test_module()) elif demisto.command() == "mongodb-write-log": return_outputs(*write_log_json()) elif demisto.command() == "mongodb-read-log": return_outputs(*read_log_json()) elif demisto.command() == "mongodb-logs-number": return_outputs(*num_log_json()) except Exception as e: return_error(f"MongoDB: {e!s}", error=e) if __name__ in ["__main__", "__builtin__", "builtins"]: main()