Add GRR approval function (#933)

* Add GRR approval function

* Add call to GRRFlowCollector

* verify access to client in SetUp

* Add more verbose download notification messages

* Fix test

* Update other tests

* Linter fix

Author: tomchop

dftimewolf/lib/collectors/grr_hosts.py

@@ -214,6 +214,37 @@ def _FindClients(self, selectors: List[str]) -> List[Client]:
         clients.append(client)
     return clients
 
+  def VerifyClientAccess(self, client: Client) -> None:
+    """Verifies and requests access to a GRR client.
+
+    This call will block until the approval is granted.
+
+    Args:
+      client: GRR client object to verify access to.
+    """
+    client_fqdn = client.data.knowledge_base.fqdn
+
+    try:
+      client.VerifyAccess()
+      self.logger.info(f"Access to {client_fqdn} granted")
+      return
+    except grr_errors.AccessForbiddenError:
+      self.logger.warning(f"No access to {client_fqdn}, requesting...")
+
+    approval = client.CreateApproval(
+      reason=self.reason,
+      notified_users=self.approvers,
+      expiration_duration_days=30,
+    )
+
+    approval_url = (
+      f"{self.grr_url}/v2/clients/{approval.client_id}/users/"
+      f"{approval.username}/approvals/{approval.approval_id}"
+    )
+    self.PublishMessage(f"Approval URL: {approval_url}")
+    approval.WaitUntilValid()
+    self.logger.info(f"Access to {client_fqdn} granted")
+
   # TODO: change object to more specific GRR type information.
   def _LaunchFlow(self, client: Client, name: str, args: str) -> str:
     """Creates the specified flow.
@@ -499,18 +530,19 @@ def _DownloadFiles(self, client: Client, flow_id: str) -> Optional[str]:
 
     flow_name = flow_handle.data.name
     if flow_name == 'TimelineFlow':
-      self.logger.debug('Downloading timeline from GRR')
+      self.logger.info('Downloading timeline from GRR')
       self._DownloadTimeline(client, flow_handle, flow_output_dir)
       return flow_output_dir
 
     if flow_name == 'OsqueryFlow':
-      self.logger.debug('Downloading osquery results from GRR')
+      self.logger.info('Downloading osquery results from GRR')
       self._DownloadOsquery(client, flow_id, flow_output_dir)
       return flow_output_dir
 
     payloads = []
     for r in flow_handle.ListResults():
       payloads.append(r.payload)
+    self.logger.info('Downloading data blobs from GRR')
     self._DownloadBlobs(client, payloads, flow_output_dir)
 
     return flow_output_dir
@@ -1424,7 +1456,11 @@ def SetUp(self,
       if host:
         client = self._GetClientBySelector(host)
         for flow_id in flows:
+          self.logger.info(
+                f'Verifying client access for {client.client_id}...'
+            )
           try:
+            self.VerifyClientAccess(client)
             client.Flow(flow_id).Get()
             self.StoreContainer(containers.GrrFlow(host, flow_id))
           except Exception as exception:  # pylint: disable=broad-except

tests/lib/collectors/grr_hosts.py

@@ -685,10 +685,17 @@ class GRRFlowCollectorTest(unittest.TestCase):
   mock_grr_api: mock.Mock
   test_state: state.DFTimewolfState
 
+  @mock.patch("grr_api_client.client.Client.VerifyAccess")
   @mock.patch('grr_api_client.flow.FlowBase.Get')
   @mock.patch('grr_api_client.client.Client.ListFlows')
   @mock.patch('grr_api_client.api.InitHttp')
-  def setUp(self, mock_InitHttp, mock_list_flows, unused_mock_flow_get):
+  def setUp(
+    self,
+    mock_InitHttp,
+    mock_list_flows,
+    unused_mock_flow_get,
+    unused_mock_verify_access,
+  ):
     self.mock_grr_api = mock.Mock()
     mock_InitHttp.return_value = self.mock_grr_api
     self.mock_grr_api.SearchClients.return_value = \
@@ -707,9 +714,10 @@ def setUp(self, mock_InitHttp, mock_list_flows, unused_mock_flow_get):
         skip_offline_clients=False,
     )
 
+  @mock.patch("grr_api_client.client.Client.VerifyAccess")
   @mock.patch('dftimewolf.lib.collectors.grr_hosts.GRRFlow._DownloadFiles')
   @mock.patch("dftimewolf.lib.collectors.grr_hosts.GRRFlow._AwaitFlow")
-  def testProcess(self, _, mock_DownloadFiles):
+  def testProcess(self, _, mock_DownloadFiles, unused_mock_verify_access):
     """Tests that the collector can be initialized."""
     self.mock_grr_api.SearchClients.return_value = \
         mock_grr_hosts.MOCK_CLIENT_LIST
@@ -730,18 +738,25 @@ def testProcess(self, _, mock_DownloadFiles):
     self.assertEqual(result.name, 'tomchop')
     self.assertEqual(result.path, '/tmp/something')
 
+  @mock.patch("grr_api_client.client.Client.VerifyAccess")
   @mock.patch('grr_api_client.flow.FlowBase.Get')
   @mock.patch('grr_api_client.client.Client.ListFlows')
   @mock.patch('grr_api_client.api.InitHttp')
   def testPreProcessNoFlows(
-    self, mock_InitHttp, mock_list_flows, unused_mock_flow_get):
+    self,
+    mock_InitHttp,
+    mock_list_flows,
+    unused_mock_flow_get,
+    unused_mock_verify_access,
+  ):
     """Tests that if no flows are found, an error is thrown."""
     self.mock_grr_api = mock.Mock()
     mock_InitHttp.return_value = self.mock_grr_api
     self.mock_grr_api.SearchClients.return_value = \
         mock_grr_hosts.MOCK_CLIENT_LIST
     mock_list_flows.return_value = [mock_grr_hosts.flow_pb_terminated]
 
+
     grr_flow_collector = grr_hosts.GRRFlowCollector(self.test_state)
     grr_flow_collector.SetUp(
         hostnames='C.0000000000000001',
@@ -762,13 +777,20 @@ def testPreProcessNoFlows(
     self.assertEqual('No flows found for collection.', error.exception.message)
     self.assertEqual(len(self.test_state.errors), 1)
 
+  @mock.patch("grr_api_client.client.Client.VerifyAccess")
   @mock.patch('grr_api_client.flow.FlowBase.Get')
   @mock.patch('grr_api_client.client.Client.ListFlows')
   @mock.patch('grr_api_client.api.InitHttp')
   @mock.patch('dftimewolf.lib.collectors.grr_hosts.GRRFlow._DownloadFiles')
   @mock.patch("dftimewolf.lib.collectors.grr_hosts.GRRFlow._AwaitFlow")
   def testProcessNoFlowData(
-    self, _, mock_DLFiles, mock_InitHttp, mock_list_flows, unused_mock_flow_get
+    self,
+    _,
+    mock_DLFiles,
+    mock_InitHttp,
+    mock_list_flows,
+    unused_mock_flow_get,
+    unused_mock_verify_access,
   ):
     """Tests Process when the flow is found but has no data collected."""
     self.mock_grr_api = mock.Mock()