Skip to content

Commit

Permalink
Add GRR approval function (#933)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
tomchop authored Nov 19, 2024
1 parent a6b44c6 commit 1b1282c
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 6 deletions.
40 changes: 38 additions & 2 deletions dftimewolf/lib/collectors/grr_hosts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
30 changes: 26 additions & 4 deletions tests/lib/collectors/grr_hosts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = \
Expand All @@ -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
Expand All @@ -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',
Expand All @@ -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()
Expand Down

0 comments on commit 1b1282c

Please sign in to comment.