def test_validate_config(app):
@app.register_parser('test1')
class TestParser:
pass
assert app.parsers
app.validate_config()
assert not app.parsers
app = Comet()
app.register_parser('test1', TestParser)
@app.register_router('test1')
def test_router(*args):
pass
app.validate_config()
assert app.parsers
def test_validate_config(app):
@app.register_parser('test1')
def parse_message(message):
pass
assert app.parsers
app.validate_config()
assert not app.parsers
app = Comet()
app.register_parser('test1', parse_message)
@app.register_router('test1')
def test_router(*args):
pass
app.validate_config()
assert app.parsers
from datetime import timedelta
from comet_common.comet_input_google_pubsub import PubSubInput
from comet_common.comet_parser_detectify import DetectifySchema
from comet_common.comet_parser_forseti import ForsetiSchema
from comet_core import Comet
from comet_core.fingerprint import comet_event_fingerprint
LOG = logging.getLogger(__name__)
LOG.setLevel(
logging.DEBUG
) # you might want to change the level to logging.INFO when running in production
LOG.info('Setting up Comet...')
APP = Comet(database_uri='sqlite:///comet-example.db'
) # use any other database backend here that SQLAlchemy supports
LOG.info('Created app instance.')
# Outputs
def send_email(recipient, subject, body):
"""Send email to recipient with given subject and body - implement your email sending or other outputs channel here.
Args:
recipient (str): recipient email address
subject (str): email subject line
body (str): email body
"""
LOG.debug(
f'Sending email to {recipient} with subject {subject}: {body[:100]}...'
def test_handle_non_addressed_events():
app = Comet()
@app.register_parser('real_time_source')
def parse_message(message):
data = json.loads(message)
return data, {}
@app.register_parser('real_time_source2')
def parse_message(message):
data = json.loads(message)
return data, {}
@app.register_config_provider('real_time_source')
def register_conf(event):
return {"escalate_cadence": timedelta(minutes=45)}
@app.register_config_provider('real_time_source2')
def register_conf(event):
return {"escalate_cadence": timedelta(minutes=45)}
app.register_real_time_source('real_time_source')
app.register_real_time_source('real_time_source2')
escalator = mock.Mock()
escalator2 = mock.Mock()
app.register_escalator('real_time_source', func=escalator)
app.register_escalator('real_time_source2', func=escalator2)
already_processed_user = '******'
# already processed real time event - needs escalation
app.data_store.add_record(
EventRecord(id=2,
received_at=datetime.utcnow() - timedelta(hours=1),
source_type='real_time_source',
owner=already_processed_user,
processed_at=datetime.utcnow() - timedelta(hours=1),
sent_at=datetime.utcnow() - timedelta(hours=1),
data={
'search_name': 'alert search name',
'name': 'needs escalation'
},
fingerprint='f2'))
# already processed real time event - still early for escalation
# the event sent 35 min ago.
app.data_store.add_record(
EventRecord(id=3,
received_at=datetime.utcnow() - timedelta(hours=1),
source_type='real_time_source2',
owner=already_processed_user,
processed_at=datetime.utcnow() - timedelta(minutes=35),
sent_at=datetime.utcnow() - timedelta(minutes=35),
data={
'search_name': 'alert search name',
'name': 'doesnt need escalation'
},
fingerprint='f3'))
app.handle_non_addressed_events()
assert escalator.call_count == 1
assert escalator2.call_count == 0
def test_process_unprocessed_whitelisted_real_time_events():
app = Comet()
app.register_parser('real_time_source', json)
app.register_real_time_source('real_time_source')
real_time_router = mock.Mock()
router = mock.Mock()
escalator = mock.Mock()
app.register_router('real_time_source', func=real_time_router)
app.register_router(func=router)
app.register_escalator(func=escalator)
check_user = '******'
# user whitelisted real time event
app.data_store.add_record(
EventRecord(id=4,
received_at=datetime.utcnow() - timedelta(days=3),
sent_at=datetime.utcnow() - timedelta(days=3),
source_type='real_time_source',
owner=check_user,
data={},
fingerprint='f4'))
app.data_store.ignore_event_fingerprint(
'f4', IgnoreFingerprintRecord.ACCEPT_RISK)
app.process_unprocessed_events()
# test the whitelisted event was not routed/escalated
assert real_time_router.call_count == 0
assert escalator.call_count == 0
def test_process_unprocessed_real_time_events():
app = Comet()
app.register_parser('real_time_source', json)
app.register_parser('datastoretest', json)
app.register_real_time_source('real_time_source')
real_time_router = mock.Mock()
router = mock.Mock()
escalator = mock.Mock()
app.register_router('real_time_source', func=real_time_router)
app.register_router(func=router)
app.register_escalator(func=escalator)
check_user = '******'
already_processed_user = '******'
# already processed regular event
app.data_store.add_record(
EventRecord(id=1,
received_at=datetime.utcnow() - timedelta(days=5),
source_type='datastoretest',
owner=already_processed_user,
data={},
processed_at=datetime.utcnow() - timedelta(days=5),
fingerprint='f1'))
# already processed real time event
app.data_store.add_record(
EventRecord(id=2,
received_at=datetime.utcnow() - timedelta(days=5),
source_type='real_time_source',
owner=already_processed_user,
processed_at=datetime.utcnow() - timedelta(days=5),
data={},
fingerprint='f2'))
# not processed real time event
app.data_store.add_record(
EventRecord(id=3,
received_at=datetime.utcnow() - timedelta(days=3),
source_type='real_time_source',
owner=check_user,
data={},
fingerprint='f3'))
# real time event needs escalation
app.data_store.add_record(
EventRecord(id=4,
received_at=datetime.utcnow() - timedelta(days=3),
sent_at=datetime.utcnow() - timedelta(days=3),
source_type='real_time_source',
owner=check_user,
data={},
fingerprint='f4'))
app.data_store.ignore_event_fingerprint(
'f4', IgnoreFingerprintRecord.ESCALATE_MANUALLY)
app.process_unprocessed_events()
assert real_time_router.call_count == 1
# route real time alert with both routers.
assert router.call_count == 1
assert real_time_router.call_args[0][2][0].owner == check_user
assert escalator.call_count == 1
def test_process_unprocessed_events():
app = Comet()
app.register_parser('datastoretest', json)
app.register_parser('datastoretest2', json)
app.register_parser('datastoretest3', json)
app.set_config('datastoretest2', {})
specific_router = mock.Mock()
router = mock.Mock()
escalator = mock.Mock()
app.register_router('datastoretest2', func=specific_router)
app.register_router(func=router)
app.register_escalator(func=escalator)
check_user = '******'
already_processed_user = '******'
app.data_store.add_record(
EventRecord(id=1,
received_at=datetime.utcnow() - timedelta(days=5),
source_type='datastoretest',
owner=already_processed_user,
data={},
processed_at=datetime.utcnow() - timedelta(days=5),
fingerprint='f1'))
app.data_store.add_record(
EventRecord(id=2,
received_at=datetime.utcnow() - timedelta(days=4),
source_type='datastoretest',
owner=already_processed_user,
data={},
fingerprint='f1'))
app.data_store.add_record(
EventRecord(id=3,
received_at=datetime.utcnow() - timedelta(days=3),
source_type='datastoretest2',
owner=check_user,
data={},
fingerprint='f3'))
app.data_store.add_record(
EventRecord(id=4,
received_at=datetime.utcnow() - timedelta(days=3),
source_type='datastoretest3',
owner=check_user,
data={},
fingerprint='f4'))
app.process_unprocessed_events()
assert specific_router.call_count == 1
assert router.call_count == 2
assert router.call_args[0][2][0].owner == check_user
assert escalator.call_count == 3
def app():
yield Comet()
def test_process_unprocessed_events_digest_mode():
app = Comet()
app.register_parser("datastoretest", json)
app.register_parser("datastoretest2", json)
app.register_parser("datastoretest3", json)
app.register_parser("datastoretest4", json)
app.set_config("datastoretest2", {})
specific_router = mock.Mock()
router = mock.Mock()
escalator = mock.Mock()
app.register_router("datastoretest2", func=specific_router)
app.register_router(func=router)
app.register_escalator(func=escalator)
check_user = "******"
already_processed_user = "******"
app.data_store.add_record(
EventRecord(
id=1,
received_at=datetime.utcnow() - timedelta(days=5),
source_type="datastoretest",
owner=already_processed_user,
data={},
processed_at=datetime.utcnow() - timedelta(days=5),
fingerprint="f1",
)
)
app.data_store.add_record(
EventRecord(
id=2,
received_at=datetime.utcnow() - timedelta(days=4),
source_type="datastoretest",
owner=already_processed_user,
data={},
fingerprint="f1",
)
)
app.data_store.add_record(
EventRecord(
id=3,
received_at=datetime.utcnow() - timedelta(days=3),
source_type="datastoretest2",
owner=check_user,
data={},
fingerprint="f3",
)
)
app.data_store.add_record(
EventRecord(
id=4,
received_at=datetime.utcnow() - timedelta(days=3),
source_type="datastoretest3",
owner=check_user,
data={},
fingerprint="f4",
)
)
app.process_unprocessed_events()
# it is expected to have specific_router called once for the datastoretest2
assert specific_router.call_count == 1
# it is expected to have two calls of the generic router for the source_type datastoretest2 and datastoretest3
assert router.call_count == 2
# and the user must be check_user
assert router.call_args[0][2][0].owner == check_user
# due to the default escalation_time=10seconds, all three events (id=2,3,4) must be escalated
assert escalator.call_count == 3
app.data_store.add_record(
EventRecord(
id=5,
received_at=datetime.utcnow() - timedelta(days=2),
source_type="datastoretest",
owner=check_user,
data={},
fingerprint="f1",
)
)
app.process_unprocessed_events()
# f1 is expected to be processed, but not sent out
assert app.data_store.get_latest_event_with_fingerprint("f1").processed_at
assert not app.data_store.get_latest_event_with_fingerprint("f1").sent_at
def test_process_unprocessed_events_non_digest_mode():
app = Comet()
app.register_parser("datastoretest4", json)
check_user = "******"
router = mock.Mock()
escalator = mock.Mock()
app.register_router(func=router)
app.register_escalator(func=escalator)
app.set_config("datastoretest4", {"communication_digest_mode": False, "new_threshold": timedelta(days=14)})
app.data_store.add_record(
EventRecord(
id=6,
received_at=datetime.utcnow() - timedelta(days=8),
source_type="datastoretest4",
sent_at=datetime.utcnow() - timedelta(days=8),
processed_at=datetime.utcnow() - timedelta(days=8),
owner=check_user,
data={},
fingerprint="f5",
)
)
app.data_store.add_record(
EventRecord(
id=7,
received_at=datetime.utcnow() - timedelta(days=2),
source_type="datastoretest4",
owner=check_user,
data={},
fingerprint="f5",
)
)
app.data_store.add_record(
EventRecord(
id=8,
received_at=datetime.utcnow(),
source_type="datastoretest4",
owner=check_user,
data={},
fingerprint="f6",
)
)
app.data_store.add_record(
EventRecord(
id=9,
received_at=datetime.utcnow() - timedelta(days=2),
source_type="datastoretest4",
sent_at=datetime.utcnow() - timedelta(days=2),
processed_at=datetime.utcnow() - timedelta(days=2),
owner=check_user,
data={},
fingerprint="f7",
)
)
app.data_store.add_record(
EventRecord(
id=10,
received_at=datetime.utcnow(),
source_type="datastoretest4",
owner=check_user,
data={},
fingerprint="f7",
)
)
# f5 is expected to be reminded
# f6 is expected to be new and sent as well
# f7 is NOT expected to be reminded
before_calling = router.call_count
app.process_unprocessed_events()
assert app.data_store.get_latest_event_with_fingerprint("f5").processed_at
assert app.data_store.get_latest_event_with_fingerprint("f6").processed_at
assert app.data_store.get_latest_event_with_fingerprint("f7").processed_at
assert router.call_count == before_calling + 1
sent_fingerprints = [e.fingerprint for e in router.call_args[0][2]]
assert "f5" in sent_fingerprints
assert "f6" in sent_fingerprints
assert "f7" not in sent_fingerprints
def app() -> Comet:
"""Returns a Comet app."""
yield Comet()
def test_handle_non_addressed_events():
app = Comet()
app.register_parser('real_time_source', json)
app.register_parser('real_time_source2', json)
app.register_real_time_source('real_time_source')
app.register_real_time_source('real_time_source2')
app.set_config('real_time_source', {'alerts': {
'alert search name':
{
'escalate_cadence': timedelta(minutes=45),
'template': 'alerts_template'
}
}})
app.set_config('real_time_source2', {'alerts': {
'alert search name':
{
'escalate_cadence': timedelta(minutes=45),
'template': 'alerts_template'
}
}})
escalator = mock.Mock()
escalator2 = mock.Mock()
app.register_escalator('real_time_source', func=escalator)
app.register_escalator('real_time_source2', func=escalator2)
already_processed_user = '******'
# already processed real time event - needs escalation
app.data_store.add_record(
EventRecord(id=2,
received_at=datetime.utcnow() - timedelta(hours=1),
source_type='real_time_source',
owner=already_processed_user,
processed_at=datetime.utcnow() - timedelta(hours=1),
sent_at=datetime.utcnow() - timedelta(hours=1),
data={'search_name': 'alert search name',
'name': 'needs escalation'},
fingerprint='f2'))
# already processed real time event - still early for escalation
# the event sent 35 min ago.
app.data_store.add_record(
EventRecord(id=3,
received_at=datetime.utcnow() - timedelta(hours=1),
source_type='real_time_source2',
owner=already_processed_user,
processed_at=datetime.utcnow() - timedelta(minutes=35),
sent_at=datetime.utcnow() - timedelta(minutes=35),
data={'search_name': 'alert search name',
'name': 'doesnt need escalation'},
fingerprint='f3'))
app.handle_non_addressed_events()
assert escalator.call_count == 1
assert escalator2.call_count == 0
