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)
