def testNoArgs(self): with mock.patch.object(grr_colab.Client, 'search', return_value=[]) as mock_fn: magics_impl.grr_search_clients_impl() mock_fn.assert_called_once_with(version=None, host=None, labels=None, mac=None, ip=None, user=None)
def testWithArgs(self): with mock.patch.object(grr_colab.Client, 'search', return_value=[]) as mock_fn: magics_impl.grr_search_clients_impl(version='test_v', host='test_host', labels=['test_label'], mac='test_mac', ip='test_ip', user='******') mock_fn.assert_called_once_with(version='test_v', host='test_host', labels=['test_label'], mac='test_mac', ip='test_ip', user='******')
def testLastSeenAgoColumn(self): current_time_secs = 1560000000 mock_clients = [ _MockClient(last_seen_at=(current_time_secs + 1) * (10**6)), _MockClient(last_seen_at=(current_time_secs - 1) * (10**6)), _MockClient(last_seen_at=(current_time_secs - 2 * 60) * (10**6)), _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)), _MockClient(last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)), ] with mock.patch.object(time, 'time', return_value=current_time_secs): with mock.patch.object(grr_colab.Client, 'search', return_value=mock_clients): df = magics_impl.grr_search_clients_impl() self.assertEqual(df.shape[0], 5) self.assertIn('last_seen_ago', df.columns) self.assertEqual(df['last_seen_ago'][0], 'in 1 seconds') self.assertEqual(df['last_seen_ago'][1], '1 seconds ago') self.assertEqual(df['last_seen_ago'][2], '2 minutes ago') self.assertEqual(df['last_seen_ago'][3], '3 hours ago') self.assertEqual(df['last_seen_ago'][4], '4 days ago')
def testOnlineColumns(self): current_time_secs = int(time.time()) mock_clients = [ _MockClient(last_seen_at=(current_time_secs - 2 * 60) * (10**6)), _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)), _MockClient(last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)), ] with mock.patch.object(grr_colab.Client, 'search', return_value=mock_clients): df = magics_impl.grr_search_clients_impl() self.assertEqual(df.shape[0], 3) self.assertIn('online', df.columns) self.assertEqual(df['online'][0], 'online') self.assertEqual(df['online'][1], 'seen-1d') self.assertEqual(df['online'][2], 'offline') self.assertIn('online.pretty', df.columns) self.assertEqual(df['online.pretty'][0], '🌕') self.assertEqual(df['online.pretty'][1], '🌓') self.assertEqual(df['online.pretty'][2], '🌑')
def testIsDataframe(self): mock_client = _MockClient() with mock.patch.object( grr_colab.Client, 'search', return_value=[mock_client]): df = magics_impl.grr_search_clients_impl() self.assertIsInstance(df, pd.DataFrame) self.assertEqual(df.shape[0], 1)
def testSortedByLastSeen(self): mock_clients = [ _MockClient(client_id='foo', last_seen_at=1000), _MockClient(client_id='bar', last_seen_at=10), _MockClient(client_id='quux', last_seen_at=100), ] with mock.patch.object( grr_colab.Client, 'search', return_value=mock_clients): df = magics_impl.grr_search_clients_impl() self.assertEqual(df.shape[0], 3) self.assertEqual(list(df['client_id']), ['foo', 'quux', 'bar'])
def testPriorityColumns(self): mock_client = _MockClient(client_id='foo') mock_client.set_hostname('test_hostname') mock_client.set_os_version('test_version') with mock.patch.object( grr_colab.Client, 'search', return_value=[mock_client]): df = magics_impl.grr_search_clients_impl() self.assertGreaterEqual(df.shape[1], 7) self.assertEqual( list(df)[:7], [ 'online.pretty', 'online', 'client_id', 'last_seen_ago', 'last_seen_at.pretty', 'knowledge_base.fqdn', 'os_info.version' ])
def grr_search_clients(line: Text) -> pd.DataFrame: """Searches for clients with specified keywords. Args: line: A string representing arguments passed to the magic command. Returns: List of clients. """ args = grr_search_clients.parser.parse_args(shlex.split(line)) return magics_impl.grr_search_clients_impl(ip=args.ip, mac=args.mac, host=args.host, user=args.user, version=args.version, labels=args.label)