def cb(self, person):
yield None
yield [Person(firstname='a', lastname='b'), Email(address='c')]
yield Person(firstname='d', lastname='d')
yield None
yield 'nope'
yield 42
yield []
yield (1, 3)
def test_output_1():
class TestCollector(BaseCollector):
config = {
'name': 'base-collector',
}
def callbacks(self):
return {Person: self.cb}
def cb(self, person):
yield None
yield [Person(firstname='a', lastname='b'), Email(address='c')]
yield Person(firstname='d', lastname='d')
yield None
yield 'nope'
yield 42
yield []
yield (1, 3)
collector = TestCollector()
result, facts = collector.collect([
Person(firstname='test-base-collector', lastname='test-base-collector')
])
assert result.executions_count == 1
assert len(result.facts) == 3
assert len(facts) == 3
def on_start(self):
print('START ENGINE')
engine_connect = kafka_producers.make_engine_connect()
engine_connect.produce(
models.EngineInfo(name='thats my engine name'),
flush=True,
)
p = Person(firstname='1', lastname='1')
p2 = Person(firstname='1', lastname='2')
e = Email(address='addr')
c = models.Collect(collector_name='toto', facts=[p, e, p2])
agent_collect = kafka_producers.make_agent_collect()
agent_collect.produce(c, flush=True)
def test_base_rate_limit():
class TestCollector(BaseCollector):
config = {
'name': 'base-collector',
'limiter': [RequestRate(limit=1, interval=Duration.SECOND)],
}
def callbacks(self):
return {Person: self.cb}
def cb(self, person):
return Person(firstname=f'res-{person.firstname}',
lastname=f'res-{person.lastname}')
collector = TestCollector()
collector.collect([Person(firstname='42', lastname='42')])
with pytest.raises(BucketFullException) as err:
collector.collect([Person(firstname='42', lastname='42')])
time.sleep(err.value.remaining_time)
collector.collect([Person(firstname='42', lastname='42')])
def test_dummy_docker_collector(mock_current_task):
mock_current_task.request.delivery_info = {'routing_key': 'dummy-docker-collector'}
res = tasks.scan.s(
[
Person(firstname='agent-test', lastname='dummy-docker-collector'),
Email(address='*****@*****.**'),
],
).apply()
result = res.get()
assert result['duration'] < 5
assert result['executions_count'] == 2
assert len(result['facts']) == 4
def test_output_raises():
class TestCollector(BaseCollector):
config = {
'name': 'base-collector',
}
def callbacks(self):
return {Person: self.cb}
def cb(self, person):
raise ValueError('aze')
collector = TestCollector()
with pytest.raises(exceptions.CollectorRuntimeError):
collector.collect([
Person(firstname='test-base-collector',
lastname='test-base-collector')
])
def test_output_3():
class TestCollector(BaseCollector):
config = {
'name': 'base-collector',
}
def callbacks(self):
return {Person: self.cb}
def cb(self, person):
pass
collector = TestCollector()
result, facts = collector.collect([
Person(firstname='test-base-collector', lastname='test-base-collector')
])
assert result.executions_count == 1
assert len(result.facts) == 0
assert facts == []
def test_output_2():
class TestCollector(BaseCollector):
config = {
'name': 'base-collector',
}
def callbacks(self):
return {Person: self.cb}
def cb(self, person):
return [Person(firstname='a', lastname='b'), Email(address='c')]
collector = TestCollector()
result, facts = collector.collect([
Person(firstname='test-base-collector', lastname='test-base-collector')
])
assert result.executions_count == 1
assert len(result.facts) == 2
assert len(facts) == 2
def cb_email(self, email):
date = self.run_container(command='date')
yield Person(firstname='dummy docker', lastname=date)
yield Email(address='dummy@email')
def cb_person(self, person):
whoami = self.run_container(command='whoami')
yield Person(firstname='dummy docker', lastname=whoami)
yield Email(address='dummy@email')
assert case['scans'] == []
# Check if the case is present in the list of cases IDs
res = api.get_all_cases()
all_cases = res['ids']
assert case_id in all_cases
# Adds a scan to our case
scan = api_models.ScanInRequest(
case_id=case_id,
scan_type='single_collector',
facts=[
api_models.FactInRequest(
fact_type='Person',
fact=Person(
firstname='test-single-collector-firstname',
lastname='test-single-collector-firstname',
),
),
],
config=models.ScanConfig(collector_name='dummy-docker-collector'),
)
res = api.create_scan(scan)
scan_id = res['id']
# Check if the scan is in the list
list_scans = api.get_all_scans()
assert scan_id in list_scans['ids']
# Check if the scan have been created
scan = api.get_scan(scan_id)
def cb(self, person):
return [Person(firstname='a', lastname='b'), Email(address='c')]
def cb(self, person):
return Person(firstname=f'res-{person.firstname}',
lastname=f'res-{person.lastname}')