def test_add_vocab():
"""Tests that vocabs can be added to a vocab list for a given field."""
vocab = Setting.find(name='vocabs')
vocab.set_vocab_for_field('Malware.type', ['banker'])
vocab.add_value_to_field_vocab('Malware.type', 'trojan')
vocablist = vocab.get_vocab_for_field('Malware.type')
assert vocablist == sorted(['banker', 'trojan'])
def test_remove_vocab():
"""Tests that a vocab can be removed from a vocab list for a given field."""
vocab = Setting.find(name='vocabs')
vocab.set_vocab_for_field('Malware.type', sorted(['banker', 'trojan']))
vocab.remove_value_from_field_vocab('Malware.type', 'trojan')
vocablist = vocab.get_vocab_for_field('Malware.type')
assert vocablist == ['banker']
def get_vocab(self, vocab):
"""Return defined vocabularies."""
try:
v = Setting.get_or_create(name='vocabs')
return v.get_vocab(vocab)
except RuntimeException as exception:
return exception, 400
def test_create_vocab_setting():
"""Tests that Vocab settings can be created."""
vocab = Setting.get_or_create(name='my_vocab', type='vocab')
vocab_ = Vocabs.get(vocab.id)
assert vocab_ is not None
assert vocab.id == vocab_.id
assert vocab_.name == 'my_vocab'
def get_killchain(self, killchain):
"""Return defined killchains."""
try:
v = Setting.get_or_create(name='killchains')
return v.get_killchain(killchain)
except RuntimeException as exception:
return exception, 400
def test_remove_vocab():
"""Tests that a vocab can be removed from a vocab list for a given field."""
vocab = Setting.find(name='malware-label-ov')
vocab.set_vocab(sorted(['banker', 'trojan']))
vocab.remove_value_from_vocab('trojan')
vocablist = vocab.get_vocab()
assert vocablist == ['banker']
def test_add_vocab():
"""Tests that vocabs can be added to a vocab list for a given field."""
vocab = Setting.find(name='malware-label-ov')
vocab.set_vocab(['banker'])
vocab.add_value_to_vocab('trojan')
vocablist = vocab.get_vocab()
assert vocablist == sorted(['banker', 'trojan'])
def test_wrong_remove_vocab_raises():
"""Tests that a RuntimeError is raised when getting a noneixstent vocab."""
vocab = Setting.find(name='vocabs')
vocab.set_vocab_for_field('Malware.type', sorted(['banker', 'trojan']))
with pytest.raises(RuntimeException):
vocab.remove_value_from_field_vocab('Notexist', ['banker'])
with pytest.raises(RuntimeException):
vocab.remove_value_from_field_vocab('Malware.type', 'notexist')
def test_create_killchain_setting():
"""Tests that killchain settings can be fetched"""
kc = Setting.get_or_create(name='my_killchain', type='killchain')
assert isinstance(kc, KillChains)
kc_ = KillChains.get(kc.id)
assert kc_ is not None
assert kc.id == kc_.id
assert kc_.name == 'my_killchain'
def set_values_for_vocab(self, vocab):
"""Set a vocabulary."""
values = request.json['values']
try:
v = Setting.get_or_create(name=vocab, type='vocab')
v.set_vocab(values)
return v.get_vocab()
except RuntimeException as exception:
return exception, 400
def get_vocab(self, vocab):
"""Return defined vocabularies."""
try:
v = Setting.find(name=vocab)
if not v:
return '', 404
return v.get_vocab()
except RuntimeException as exception:
return exception, 400
def set_killchain_pahses(self, killchain):
"""Set phases for a killchain."""
phases = request.json['phases']
try:
kc = Setting.get_or_create(name=killchain)
kc.set_phases(phases)
return kc.get_killchain()
except RuntimeException as exception:
return exception, 400
def add_phase_to_killchain(self, killchain):
"""Set phases for a killchain."""
phase = request.json['phase']
try:
kc = Setting.get_or_create(name=killchain)
kc.add_phase_to_killchain(phase)
return kc.get_killchain()
except RuntimeException as exception:
return exception, 400
def remove_value_from_field_vocab(self, field):
"""Return defined vocabularies for a field."""
value = request.json['value']
v = Setting.find(name='vocabs')
try:
v.remove_value_from_field_vocab(field, value)
except RuntimeException as exception:
return exception, 400
return v.get_vocab_for_field(field)
def get_killchain(self, killchain):
"""Return defined killchains."""
try:
kc = Setting.find(name=killchain)
if not kc:
return '', 404
return kc.get_killchain()
except RuntimeException as exception:
return exception, 400
def add_phase_to_killchain(self, killchain):
"""Set phases for a killchain."""
phase = request.json['phase']
try:
kc = Setting.find(name='killchains')
kc.add_phase_to_killchain(killchain, phase)
return kc.get_killchain(killchain)
except RuntimeException as exception:
return exception, 400
def remove_phase_from_killchain(self, killchain):
"""Remove a phase from a killchain."""
phase = request.json['phase']
kc = Setting.find(name='killchains')
try:
kc.remove_phase_from_killchain(killchain, phase)
except RuntimeException as exception:
return exception, 400
return kc.get_killchain(killchain)
def remove_value_from_vocab(self, vocab):
"""Remove a value from a vocab."""
value = request.json['value']
v = Setting.find(name='vocabs')
try:
v.remove_value_from_vocab(vocab, value)
except RuntimeException as exception:
return exception, 400
return v.get_vocab(vocab)
def add_value_to_vocab(self, vocab):
"""Set a vocabulary."""
value = request.json['value']
try:
v = Setting.find(name='vocabs')
v.add_value_to_vocab(vocab, value)
return v.get_vocab(vocab)
except RuntimeException as exception:
return exception, 400
def populate_settings():
v = Setting(name='malware-label-ov', type='vocab').save()
v.set_vocab(sorted(['adware', 'backdoor']))
return [v]
def test_filter_vocab():
"""Tests that a vocab list for a given list can be filtered."""
vocab = Setting.find(name='malware-label-ov')
vocab.set_vocab(sorted(['banker', 'trojan']))
assert vocab.filter_values_vocab('tro') == ['trojan']
def test_set_vocab():
"""Tests that a vocab can be set for a given field."""
vocab = Setting.find(name='vocabs')
vocab.set_vocab('malware-label-ov', ['toto'])
vocablist = vocab.get_vocab('malware-label-ov')
assert vocablist == ['toto']
def test_wrong_field_raises():
"""Tests that a RuntimeError is raised when getting a noneixstent vocab."""
vocab = Setting.find(name='vocabs')
with pytest.raises(RuntimeException):
vocab.get_vocab_for_field('Notexist')
def get_vocab_for_field(self, field):
"""Return defined vocabularies for a field."""
v = Setting.find(name='vocabs')
return v.get_vocab_for_field(field)
def test_filter_vocab():
"""Tests that a vocab list for a given list can be filtered."""
vocab = Setting.find(name='vocabs')
vocab.set_vocab_for_field('Malware.type', sorted(['banker', 'trojan']))
assert vocab.filter_values_for_field_vocab('Malware.type',
'tro') == ['trojan']
def add_value_to_field_vocab(self, field):
"""Return defined vocabularies for a field."""
value = request.json['value']
v = Setting.find(name='vocabs')
v.add_value_to_field_vocab(field, value)
return v.get_vocab_for_field(field)
def get_killchains(self):
"""Return all killchains"""
try:
return Setting.filter({'type': 'killchain'})
except RuntimeException as exception:
return exception, 400
def test_set_vocab():
"""Tests that a vocab can be set for a given field."""
vocab = Setting.find(name='vocabs')
vocab.set_vocab_for_field('Malware.type', ['banker', 'other'])
vocablist = vocab.get_vocab_for_field('Malware.type')
assert vocablist == ['banker', 'other']
def test_get_vocab_setting():
"""Tests that a vocab setting can be saved and has correct type."""
vocab = Setting.find(name='vocabs')
assert isinstance(vocab, Vocabs)
