DR-496 Add PAM Discovery

Add to the pam action commands the ability to start a discovery job.

keepercommander/commands/discover/__init__.py

@@ -0,0 +1,245 @@
+from __future__ import annotations
+import logging
+from ..base import Command
+from ..pam.config_facades import PamConfigurationRecordFacade
+from ..pam import gateway_helper
+from ..pam.router_helper import get_response_payload
+from ..pam.gateway_helper import get_all_gateways
+from ..ksm import KSMCommand
+from ... import utils, vault_extensions
+from ... import vault
+from ...proto import APIRequest_pb2
+from ...crypto import encrypt_aes_v2, decrypt_aes_v2
+from ...display import bcolors
+from discovery_common.constants import PAM_USER, PAM_MACHINE, PAM_DATABASE, PAM_DIRECTORY
+import json
+import base64
+
+from typing import List, Optional, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ...params import KeeperParams
+    from ...vault import KeeperRecord, ApplicationRecord
+    from ...proto import pam_pb2
+
+
+class GatewayContext:
+    def __init__(self, configuration: KeeperRecord, facade: PamConfigurationRecordFacade,
+                 gateway: pam_pb2.PAMController, application: ApplicationRecord):
+        self.configuration = configuration
+        self.facade = facade
+        self.gateway = gateway
+        self.application = application
+        self._shared_folders = None
+
+    @staticmethod
+    def from_configuration_uid(params: KeeperParams, configuration_uid: str):
+
+        configuration_record = vault.KeeperRecord.load(params, configuration_uid)
+        if not isinstance(configuration_record, vault.TypedRecord):
+            print(f'{bcolors.FAIL}PAM Configuration [{configuration_uid}] is not available.{bcolors.ENDC}')
+            return
+
+        configuration_facade = PamConfigurationRecordFacade()
+        configuration_facade.record = configuration_record
+
+        gateway_uid = configuration_facade.controller_uid
+        gateway = next((x for x in gateway_helper.get_all_gateways(params)
+                        if utils.base64_url_encode(x.controllerUid) == gateway_uid),
+                       None)
+
+        if gateway is None:
+            print(f'{bcolors.FAIL}Gateway [{gateway_uid}] was not found.{bcolors.ENDC}')
+            return
+
+        application_id = utils.base64_url_encode(gateway.applicationUid)
+        application = KSMCommand.get_app_record(params, application_id)
+
+        return GatewayContext(
+            configuration=configuration_record,
+            facade=configuration_facade,
+            gateway=gateway,
+            application=application
+        )
+
+    @staticmethod
+    def from_gateway(params: KeeperParams, gateway: str):
+        # Get all the PAM configuration records
+        configuration_records = list(vault_extensions.find_records(params, "pam.*Configuration"))
+        if len(configuration_records) == 0:
+            print(f"{bcolors.FAIL}Cannot find any PAM configuration records in the Vault{bcolors.ENDC}")
+
+        all_gateways = get_all_gateways(params)
+
+        for record in configuration_records:
+
+            logging.debug(f"checking configuration record {record.title}")
+
+            # Load the configuration record and get the gateway_uid from the facade.
+            configuration_record = vault.KeeperRecord.load(params, record.record_uid)
+            configuration_facade = PamConfigurationRecordFacade()
+            configuration_facade.record = configuration_record
+
+            configuration_gateway_uid = configuration_facade.controller_uid
+            if configuration_gateway_uid is None:
+                logging.debug(f"configuration {configuration_record.title} does not have a gateway set, skipping.")
+                continue
+
+            # Get the gateway for this configuration
+            found_gateway = next((x for x in all_gateways if utils.base64_url_encode(x.controllerUid) ==
+                                  configuration_gateway_uid), None)
+            if found_gateway is None:
+                logging.debug(f"cannot find gateway for configuration {configuration_record.title}, skipping.")
+                continue
+
+            application_id = utils.base64_url_encode(found_gateway.applicationUid)
+            application = KSMCommand.get_app_record(params, application_id)
+            if application is None:
+                logging.debug(f"cannot find application for gateway {gateway}, skipping.")
+
+            if (utils.base64_url_encode(found_gateway.controllerUid) == gateway or
+                    found_gateway.controllerName.lower() == gateway.lower()):
+                return GatewayContext(
+                    configuration=configuration_record,
+                    facade=configuration_facade,
+                    gateway=found_gateway,
+                    application=application
+                )
+
+        return None
+
+    @property
+    def gateway_uid(self) -> str:
+        return utils.base64_url_encode(self.gateway.controllerUid)
+
+    @property
+    def configuration_uid(self) -> str:
+        return self.configuration.record_uid
+
+    @property
+    def gateway_name(self) -> str:
+        return self.gateway.controllerName
+
+    @property
+    def default_shared_folder_uid(self) -> str:
+        return self.facade.folder_uid
+
+    def is_gateway(self, request_gateway: str) -> bool:
+        if request_gateway is None or self.gateway_name is None:
+            return False
+        return (request_gateway == utils.base64_url_encode(self.gateway.controllerUid) or
+                request_gateway.lower() == self.gateway_name.lower())
+
+    def get_shared_folders(self, params: KeeperParams) -> List[dict]:
+        if self._shared_folders is None:
+            self._shared_folders = []
+            application_uid = utils.base64_url_encode(self.gateway.applicationUid)
+            app_info = KSMCommand.get_app_info(params, application_uid)
+            for info in app_info:
+                if info.shares is None:
+                    continue
+                for shared in info.shares:
+                    uid_str = utils.base64_url_encode(shared.secretUid)
+                    shared_type = APIRequest_pb2.ApplicationShareType.Name(shared.shareType)
+                    if shared_type == 'SHARE_TYPE_FOLDER':
+                        if uid_str not in params.shared_folder_cache:
+                            continue
+                        cached_shared_folder = params.shared_folder_cache[uid_str]
+                        self._shared_folders.append({
+                            "uid": uid_str,
+                            "name": cached_shared_folder.get('name_unencrypted'),
+                            "folder": cached_shared_folder
+                        })
+        return self._shared_folders
+
+    def decrypt(self, cipher_base64: bytes) -> dict:
+        ciphertext = base64.b64decode(cipher_base64)
+        return json.loads(decrypt_aes_v2(ciphertext, self.configuration.record_key))
+
+    def encrypt(self, data: dict) -> str:
+        json_data = json.dumps(data)
+        ciphertext = encrypt_aes_v2(json_data.encode(), self.configuration.record_key)
+        return base64.b64encode(ciphertext).decode()
+
+    def encrypt_str(self, data: Union[bytes, str]) -> str:
+        if isinstance(data, str):
+            data = data.encode()
+        ciphertext = encrypt_aes_v2(data, self.configuration.record_key)
+        return base64.b64encode(ciphertext).decode()
+
+
+class PAMGatewayActionDiscoverCommandBase(Command):
+
+    """
+    The discover command base.
+
+    Contains static methods to get the configuration record, get and update the discovery store. These are methods
+    used by multiple discover actions.
+    """
+
+    # If the discovery data field does not exist, or the field contains no values, use the template to init the
+    # field.
+
+    STORE_LABEL = "discoveryKey"
+    FIELD_MAPPING = {
+        "hostnameOrIPAddress": {
+            "type": "dict",
+            "field_input": [
+                {"key": "hostName", "prompt": "Hostname"},
+                {"key": "port", "prompt": "Port"}
+            ],
+            "field_format": [
+                {"key": "hostName", "label": "Hostname"},
+                {"key": "port", "label": "Port"},
+            ]
+        },
+        "alternativeIPs": {
+            "type": "multiline",
+        }
+    }
+
+    type_name_map = {
+        PAM_USER: "PAM Users",
+        PAM_MACHINE: "PAM Machines",
+        PAM_DATABASE: "PAM Databases",
+        PAM_DIRECTORY: "PAM Directories",
+    }
+
+    @staticmethod
+    def get_response_data(router_response: dict) -> Optional[dict]:
+
+        if router_response is None:
+            return None
+
+        response = router_response.get("response")
+        logging.debug(f"Router Response: {response}")
+        payload = get_response_payload(router_response)
+        return payload.get("data")
+
+    @staticmethod
+    def _gr(msg):
+        return f"{bcolors.OKGREEN}{msg}{bcolors.ENDC}"
+
+    @staticmethod
+    def _bl(msg):
+        return f"{bcolors.OKBLUE}{msg}{bcolors.ENDC}"
+
+    @staticmethod
+    def _h(msg):
+        return f"{bcolors.HEADER}{msg}{bcolors.ENDC}"
+
+    @staticmethod
+    def _b(msg):
+        return f"{bcolors.BOLD}{msg}{bcolors.ENDC}"
+
+    @staticmethod
+    def _f(msg):
+        return f"{bcolors.FAIL}{msg}{bcolors.ENDC}"
+
+    @staticmethod
+    def _p(msg):
+        return msg
+
+    @staticmethod
+    def _n(record_type):
+        return PAMGatewayActionDiscoverCommandBase.type_name_map.get(record_type, "PAM Configuration")