Add support for rotate on expire to share folder for pam user type records

Author: mtyagi-ks

keepercli-package/src/keepercli/commands/pam/discovery/__init__.py

@@ -200,7 +200,10 @@ def from_gateway(vault: vault_online.VaultOnline, gateway: str, configuration_ui
             application_id = utils.base64_url_encode(found_gateway.applicationUid)
             application = vault.vault_data.load_record(application_id)
             if application is None:
-                logger.debug(f"cannot find application for gateway {gateway}, skipping.")
+                logger.warning(
+                    f"KSM application for gateway {gateway} is not in the vault "
+                    f"(record {application_id}); discovery may still work via the router."
+                )
 
             if (utils.base64_url_encode(found_gateway.controllerUid) == gateway or
                     found_gateway.controllerName.lower() == gateway.lower()):

keepercli-package/src/keepercli/commands/pam/discovery/discover.py

@@ -18,7 +18,7 @@
 
 from keepersdk.helpers.pam_user_record_facade import PamUserRecordFacade
 from keepersdk.helpers.keeper_dag.jobs import Jobs
-from keepersdk.helpers.keeper_dag.dag_types import (CredentialBase, DiscoveryDelta, DiscoveryObject, JobItem, UserAcl, DirectoryInfo, 
+from keepersdk.helpers.keeper_dag.dag_types import (CredentialBase, DiscoveryDelta, DiscoveryObject, JobItem, Settings, UserAcl, DirectoryInfo, 
                 BulkRecordConvert, BulkRecordAdd, BulkRecordSuccess, BulkProcessResults, NormalizedRecord, BulkRecordFail, PromptResult,
                 PromptActionEnum, RecordField)
 from keepersdk.helpers.keeper_dag.dag_vertex import DAGVertex
@@ -152,7 +152,7 @@ def print_job_detail(vault: vault_online.VaultOnline,
                          job_id: str):
 
         def _find_job(configuration_record) -> Optional[Dict]:
-            jobs_obj = Jobs(record=configuration_record)
+            jobs_obj = Jobs(record=configuration_record, vault=vault)
             job_item = jobs_obj.get_job(job_id)
             if job_item is not None:
                 return {
@@ -167,7 +167,7 @@ def _find_job(configuration_record) -> Optional[Dict]:
         if gateway_context is not None:
             jobs = payload["jobs"]
             job = jobs.get_job(job_id)
-            infra = Infrastructure(record=gateway_context.configuration)
+            infra = Infrastructure(record=gateway_context.configuration, vault=vault)
 
             status = "RUNNING"
             if job.end_ts is not None and not job.error:
@@ -296,7 +296,7 @@ def execute(self, context: KeeperParams, **kwargs):
                 if len(gateway_context.gateway_name) > max_gateway_name:
                     max_gateway_name = len(gateway_context.gateway_name)
 
-                jobs = Jobs(record=configuration_record)
+                jobs = Jobs(record=configuration_record, vault=vault)
                 if show_history is True:
                     job_list = reversed(jobs.history)
                 else:
@@ -391,7 +391,7 @@ def execute(self, context: KeeperParams, **kwargs):
             multi_conf_msg(gateway, err)
             return
 
-        jobs = Jobs(record=gateway_context.configuration)
+        jobs = Jobs(record=gateway_context.configuration, vault=vault)
         current_job_item = jobs.current_job
         removed_prior_job = None
         if current_job_item is not None:
@@ -467,15 +467,20 @@ def execute(self, context: KeeperParams, **kwargs):
                         setattr(c, key, obj[key])
                     credentials.append(c.model_dump())
 
+        user_map_entries = self.make_protobuf_user_map(
+            context=context,
+            gateway_context=gateway_context
+        )
+        if len(user_map_entries) == 0:
+            logger.info(
+                "No pamUser records are linked to this configuration; "
+                "discovery will run without an existing user map."
+            )
+
         action_inputs = GatewayActionDiscoverJobStartInputs(
             configuration_uid=gateway_context.configuration_uid,
             resource_uid=kwargs.get('resource_uid'),
-            user_map=gateway_context.encrypt(
-                self.make_protobuf_user_map(
-                    context=context,
-                    gateway_context=gateway_context
-                )[0]
-            ),
+            user_map=gateway_context.encrypt({"users": user_map_entries}),
 
             shared_folder_uid=gateway_context.default_shared_folder_uid,
             languages=[kwargs.get('language')],
@@ -507,16 +512,39 @@ def execute(self, context: KeeperParams, **kwargs):
             logger.error(f"The router returned a failure.")
             return
 
+        discovery_settings = Settings(
+            credentials=[CredentialBase(**c) for c in credentials],
+            default_shared_folder_uid=gateway_context.default_shared_folder_uid,
+            include_azure_aadds=kwargs.get('include_azure_aadds', False),
+            skip_rules=kwargs.get('skip_rules', False),
+            skip_machines=kwargs.get('skip_machines', False),
+            skip_databases=kwargs.get('skip_databases', False),
+            skip_directories=kwargs.get('skip_directories', False),
+            skip_cloud_users=kwargs.get('skip_cloud_users', False),
+            user_map=user_map_entries or None,
+        )
+        job_id = jobs.start(
+            settings=discovery_settings,
+            resource_uid=kwargs.get('resource_uid'),
+            conversation_id=conversation_id,
+        )
+        jobs.close()
+
         if "has been queued" in data.get("Response", ""):
 
             if removed_prior_job is None:
-                logger.info("The discovery job is currently running.")
+                logger.info(f"Discovery job {job_id} is running.")
             else:
-                logger.info(f"Active discovery job {removed_prior_job} has been removed and new discovery job is running.")
+                logger.info(
+                    f"Active discovery job {removed_prior_job} has been removed; "
+                    f"discovery job {job_id} is running."
+                )
             logger.info(f"To check the status, use the command 'pam action discover status'.")
-            logger.info(f"To stop and remove the current job, use the command 'pam action discover remove -j <Job ID>'.")
+            logger.info(f"To stop and remove the current job, use the command 'pam action discover remove -j {job_id}'.")
         else:
             router_utils.print_router_response(router_response, "job_info", conversation_id, gateway_uid=gateway_context.gateway_uid)
+            logger.info(f"Discovery job {job_id} was recorded on the configuration.")
+            logger.info(f"To check the status, use the command 'pam action discover status -j {job_id}'.")
 
     @staticmethod
     def make_protobuf_user_map(context: KeeperParams, gateway_context: GatewayContext) -> List[dict]:
@@ -580,7 +608,7 @@ def execute(self, context: KeeperParams, **kwargs):
         all_gateways = GatewayContext.all_gateways(vault)
 
         def _find_job(configuration_record) -> Optional[Dict]:
-            jobs_obj = Jobs(record=configuration_record)
+            jobs_obj = Jobs(record=configuration_record, vault=vault)
             job_item = jobs_obj.get_job(job_id)
             if job_item is not None:
                 return {
@@ -1775,7 +1803,7 @@ def _get_directory_info(domain: str,
     def remove_job(context: KeeperParams, configuration_record: vault_record.KeeperRecord, job_id: str):
 
         try:
-            jobs = Jobs(record=configuration_record, context=context)
+            jobs = Jobs(record=configuration_record, vault=context.vault)
             jobs.cancel(job_id)
             logger.info(f"No items left to process. Removing completed discovery job.")
         except Exception as err:
@@ -1786,7 +1814,7 @@ def preview(self, job_item: JobItem, context: KeeperParams, gateway_context: Gat
 
         sync_point = job_item.sync_point
         infra = Infrastructure(record=gateway_context.configuration,
-                               context=context,
+                               vault=context.vault,
                                logger=logger,
                                debug_level=debug_level)
         infra.load(sync_point)
@@ -1941,7 +1969,7 @@ def execute(self, context: KeeperParams, **kwargs):
 
             # Get the current job.
             # There can only be one active job.
-            jobs = Jobs(record=configuration_record, context=context, logger=logger, debug_level=debug_level)
+            jobs = Jobs(record=configuration_record, vault=vault, logger=logger, debug_level=debug_level)
             job_item = jobs.current_job
             if job_item is None:
                 continue

keepercli-package/src/keepercli/commands/pam/pam_gateway_action.py

@@ -255,7 +255,10 @@ def record_rotate(self, context: KeeperParams, record_uid, slient:bool = False):
             config_uid = facade.controller_uid
 
         if not resource_uid:
-            tmp_dag = tunnel_graph.TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, record.record_uid)
+            tmp_dag = tunnel_graph.TunnelDAG(
+                vault, encrypted_session_token, encrypted_transmission_key, record.record_uid,
+                transmission_key=transmission_key,
+            )
             resource_uid = tmp_dag.get_resource_uid(record_uid)
             if not resource_uid:
                 is_noop = False

keepercli-package/src/keepercli/commands/pam/pam_rotation.py

@@ -295,7 +295,8 @@ def execute(self, context: KeeperParams, **kwargs):
         def config_resource(_dag, target_record, target_config_uid, silent=None):
             if not _dag.linking_dag.has_graph:
                 if target_config_uid:
-                    _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_config_uid)
+                    _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_config_uid,
+                                     transmission_key=transmission_key)
                     _dag.edit_tunneling_config(rotation=True)
                 else:
                     raise base.CommandError(f'Resource "{target_record.record_uid}" is not associated '
@@ -305,7 +306,7 @@ def config_resource(_dag, target_record, target_config_uid, silent=None):
             resource_dag = None
             if not _dag.resource_belongs_to_config(target_record.record_uid):
                 resource_dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key,
-                                         target_record.record_uid)
+                                         target_record.record_uid, transmission_key=transmission_key)
                 _dag.link_resource_to_config(target_record.record_uid)
 
             admin = kwargs.get('admin')
@@ -401,10 +402,12 @@ def config_iam_aad_user(_dag, target_record, target_iam_aad_config_uid):
                 return
 
             if _dag and not _dag.linking_dag.has_graph:
-                _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_iam_aad_config_uid)
+                _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_iam_aad_config_uid,
+                                 transmission_key=transmission_key)
                 if not _dag or not _dag.linking_dag.has_graph:
                     _dag.edit_tunneling_config(rotation=True)
-            old_dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_record.record_uid)
+            old_dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_record.record_uid,
+                                transmission_key=transmission_key)
             if old_dag.linking_dag.has_graph and old_dag.record.record_uid != target_iam_aad_config_uid:
                 old_dag.remove_from_dag(target_record.record_uid)
 
@@ -621,7 +624,8 @@ def config_user(_dag, target_record, target_resource_uid, target_config_uid=None
                     return
 
             if isinstance(target_resource_uid, str) and len(target_resource_uid) > 0:
-                _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_resource_uid)
+                _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_resource_uid,
+                                   transmission_key=transmission_key)
                 if not _dag or not _dag.linking_dag.has_graph:
                     if target_config_uid and target_resource_uid:
                         config_resource(_dag, target_record, target_config_uid, silent=silent)
@@ -639,7 +643,8 @@ def config_user(_dag, target_record, target_resource_uid, target_config_uid=None
             current_record_rotation = context.get_record_rotation(target_record.record_uid)
 
             if not _dag or not _dag.linking_dag.has_graph:
-                _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_resource_uid)
+                _dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, target_resource_uid,
+                                   transmission_key=transmission_key)
                 if not _dag.linking_dag.has_graph:
                     raise base.CommandError(f'Resource "{target_resource_uid}" is not associated '
                                            f'with any configuration. '
@@ -824,6 +829,8 @@ def config_user(_dag, target_record, target_resource_uid, target_config_uid=None
         if record_name:
             if record_name in vault.vault_data._records:
                 record_uids.add(record_name)
+            elif vault.vault_data.load_record(record_name):
+                record_uids.add(record_name)
             else:
                 rs = folder_utils.try_resolve_path(context, record_name)
                 if rs is not None:
@@ -866,7 +873,10 @@ def add_folders(folder: vault_types.Folder):
         if folder_uids:
             regex = re.compile(fnmatch.translate(record_pattern), re.IGNORECASE).match if record_pattern else None
             for folder_uid in folder_uids:
-                folder_records = vault.vault_data.get_folder(folder_uid).records
+                folder = vault.vault_data.get_folder(folder_uid)
+                if not folder:
+                    continue
+                folder_records = folder.records
                 if not folder_records:
                     continue
                 if record_pattern and record_pattern in folder_records:
@@ -957,7 +967,8 @@ def add_folders(folder: vault_types.Folder):
         r_requests = []
 
         for _record in pam_records:
-            tmp_dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, _record.record_uid)
+            tmp_dag = TunnelDAG(vault, encrypted_session_token, encrypted_transmission_key, _record.record_uid,
+                                 transmission_key=transmission_key)
             if _record.record_type in ['pamMachine', 'pamDatabase', 'pamDirectory', 'pamRemoteBrowser']:
                 config_resource(tmp_dag, _record, config_uid, silent=kwargs.get('silent'))
             elif _record.record_type == 'pamUser':
@@ -1108,7 +1119,7 @@ def is_resource_ok(resource_id, vault, configuration_uid):
             logger.info(f"Is Rotation Disabled: {rri.disabled}")
 
             rq = pam_pb2.PAMGenericUidsRequest()
-            schedules_proto = router_utils.router_get_rotation_schedules(context, rq)
+            schedules_proto = router_utils.router_get_rotation_schedules(vault, rq)
             if schedules_proto:
                 schedules = list(schedules_proto.schedules)
                 for s in schedules: