class TestDatabase:
def setUp(self):
self.d = Database(dsn="sqlite://")
def test_machine_add_clean(self):
# Add.
self.d.add_machine("a", "a", "1.1.1.1", "win", "", "", "", "", "")
session = self.d.Session()
assert_equal(session.query(Machine).count(), 1)
# Delete.
self.d.clean_machines()
assert_equal(session.query(Machine).count(), 0)
def test_task_add_del(self):
# Add.
sample_path = tempfile.mkstemp()[1]
self.d.add_path(sample_path)
session = self.d.Session()
assert_equal(session.query(Sample).count(), 1)
assert_equal(session.query(Task).count(), 1)
# Drop tasks.
self.d.drop_tasks()
assert_equal(session.query(Task).count(), 0)
assert_equal(session.query(Sample).count(), 1)
# Drop samples.
self.d.drop_samples()
assert_equal(session.query(Sample).count(), 0)
def cuckoo_dedup_cluster_queue():
"""
Cleans duplicated pending tasks from cluster queue
"""
main_db = Database()
session = main_db.Session()
dist_session = create_session(rep_config.distributed.db, echo=False)
dist_db = dist_session()
hash_dict = dict()
duplicated = session.query(Sample, Task).join(Task).filter(
Sample.id == Task.sample_id,
Task.status == "pending").order_by(Sample.sha256)
for sample, task in duplicated:
try:
# hash -> [[id, file]]
hash_dict.setdefault(sample.sha256, list())
hash_dict[sample.sha256].append((task.id, task.target))
except UnicodeDecodeError:
pass
resolver_pool.map(
lambda sha256: dist_delete_data(hash_dict[sha256][1:], dist_db),
hash_dict)
class TestDropDatabase:
"""Tests database creation, adding a couple of tasks and dropping the db."""
def setUp(self):
self.d = Database(dsn="sqlite://")
# When dropping a db, any other query raises a (OperationalError) no such table.
@raises(OperationalError)
def test_drop(self):
# Add task.
sample_path = tempfile.mkstemp()[1]
self.d.add_path(sample_path)
session = self.d.Session()
assert_equal(session.query(Sample).count(), 1)
assert_equal(session.query(Task).count(), 1)
# Add url.
self.d.add_url("http://foo.bar")
assert_equal(session.query(Sample).count(), 1)
assert_equal(session.query(Task).count(), 2)
# Drop.
self.d.drop()
assert_equal(session.query(Task).count(), 0)
class TestDatabaseEngine(object):
"""Test database stuff."""
URI = None
def setup_method(self):
self.temp = tempfile.mktemp()
with open(self.temp, "w") as f:
f.write("hehe")
self.d = Database(dsn="sqlite://")
# self.d.connect(dsn=self.URI)
self.session = self.d.Session()
def teardown_method(self):
del self.d
os.unlink(self.temp)
def add_url(self, url, priority=1, status="pending"):
task_id = self.d.add_url(url, priority=priority)
self.d.set_status(task_id, status)
return task_id
def test_add_tasks(self):
# Add task.
count = self.session.query(Task).count()
self.d.add_path(self.temp)
assert self.session.query(Task).count() == count + 1
# Add url.
self.d.add_url("http://foo.bar")
assert self.session.query(Task).count() == count + 2
def test_error_exists(self):
task_id = self.add_url("http://google.com/")
self.d.add_error("A" * 1024, task_id)
assert len(self.d.view_errors(task_id)) == 1
self.d.add_error("A" * 1024, task_id)
assert len(self.d.view_errors(task_id)) == 2
def test_long_error(self):
self.add_url("http://google.com/")
self.d.add_error("A" * 1024, 1)
err = self.d.view_errors(1)
assert err and len(err[0].message) == 1024
def test_task_set_options(self):
assert self.d.add_path(self.temp, options={"foo": "bar"}) is None
t1 = self.d.add_path(self.temp, options="foo=bar")
assert self.d.view_task(t1).options == "foo=bar"
def test_task_tags_str(self):
task = self.d.add_path(self.temp, tags="foo,,bar")
tag_list = list(self.d.view_task(task).tags)
assert [str(x.name) for x in tag_list].sort() == ["foo", "bar"].sort()
def test_add_machine(self):
self.d.add_machine(
name="name1",
label="label",
ip="1.2.3.4",
platform="windows",
tags="tag1 tag2",
interface="int0",
snapshot="snap0",
resultserver_ip="5.6.7.8",
resultserver_port=2043,
)
self.d.add_machine(
name="name2",
label="label",
ip="1.2.3.4",
platform="windows",
tags="tag1 tag2",
interface="int0",
snapshot="snap0",
resultserver_ip="5.6.7.8",
resultserver_port=2043,
)
m1 = self.d.view_machine("name1")
m2 = self.d.view_machine("name2")
assert m1.to_dict() == {
"status": None,
"locked": False,
"name": u"name1",
"resultserver_ip": u"5.6.7.8",
"ip": u"1.2.3.4",
"tags": [u"tag1tag2"],
"label": u"label",
"locked_changed_on": None,
"platform": u"windows",
"snapshot": u"snap0",
"interface": u"int0",
"status_changed_on": None,
"id": 1,
"resultserver_port": u"2043",
}
assert m2.to_dict() == {
"id": 2,
"interface": u"int0",
"ip": u"1.2.3.4",
"label": u"label",
"locked": False,
"locked_changed_on": None,
"name": u"name2",
"platform": u"windows",
"resultserver_ip": u"5.6.7.8",
"resultserver_port": u"2043",
"snapshot": u"snap0",
"status": None,
"status_changed_on": None,
"tags": [u"tag1tag2"],
}
class TestDatabaseEngine:
"""Test database stuff."""
URI = None
def setup_method(self):
with NamedTemporaryFile(mode="w+", delete=False) as f:
f.write("hehe")
self.temp_filename = f.name
pcap_header_base64 = b"1MOyoQIABAAAAAAAAAAAAAAABAABAAAA"
pcap_bytes = base64.b64decode(pcap_header_base64)
self.temp_pcap = store_temp_file(pcap_bytes, "%s.pcap" % f.name)
self.d = Database(dsn="sqlite://")
# self.d.connect(dsn=self.URI)
self.session = self.d.Session()
self.binary_storage = os.path.join(os.getcwd(), "storage/binaries")
os.makedirs(self.binary_storage)
def teardown_method(self):
del self.d
os.unlink(self.temp_filename)
shutil.rmtree(self.binary_storage)
def add_url(self, url, priority=1, status="pending"):
task_id = self.d.add_url(url, priority=priority)
self.d.set_status(task_id, status)
return task_id
def test_add_tasks(self):
# Add task.
count = self.session.query(Task).count()
self.d.add_path(self.temp_filename)
assert self.session.query(Task).count() == count + 1
# Add url.
self.d.add_url("http://foo.bar")
assert self.session.query(Task).count() == count + 2
def test_error_exists(self):
task_id = self.add_url("http://google.com/")
self.d.add_error("A" * 1024, task_id)
assert len(self.d.view_errors(task_id)) == 1
self.d.add_error("A" * 1024, task_id)
assert len(self.d.view_errors(task_id)) == 2
def test_long_error(self):
self.add_url("http://google.com/")
self.d.add_error("A" * 1024, 1)
err = self.d.view_errors(1)
assert err and len(err[0].message) == 1024
def test_task_set_options(self):
assert self.d.add_path(self.temp_filename, options={"foo": "bar"
}) is None
t1 = self.d.add_path(self.temp_filename, options="foo=bar")
assert self.d.view_task(t1).options == "foo=bar"
def test_task_tags_str(self):
task = self.d.add_path(self.temp_filename, tags="foo,,bar")
tag_list = list(self.d.view_task(task).tags)
assert [str(x.name) for x in tag_list].sort() == ["foo", "bar"].sort()
def test_reschedule_file(self):
count = self.session.query(Task).count()
task_id = self.d.add_path(self.temp_filename)
assert self.session.query(Task).count() == count + 1
task = self.d.view_task(task_id)
assert task is not None
assert task.category == "file"
# write a real sample to storage
sample_path = os.path.join(self.binary_storage, task.sample.sha256)
shutil.copy(self.temp_filename, sample_path)
new_task_id = self.d.reschedule(task_id)
assert new_task_id is not None
new_task = self.d.view_task(new_task_id)
assert new_task.category == "file"
def test_reschedule_static(self):
count = self.session.query(Task).count()
task_id = self.d.add_static(self.temp_filename)
assert self.session.query(Task).count() == count + 1
task = self.d.view_task(task_id)
assert task is not None
assert task.category == "static"
# write a real sample to storage
static_path = os.path.join(self.binary_storage, task.sample.sha256)
shutil.copy(self.temp_filename, static_path)
new_task_id = self.d.reschedule(task_id)
assert new_task_id is not None
new_task = self.d.view_task(new_task_id)
assert new_task.category == "static"
def test_reschedule_pcap(self):
count = self.session.query(Task).count()
task_id = self.d.add_pcap(self.temp_pcap)
assert self.session.query(Task).count() == count + 1
task = self.d.view_task(task_id)
assert task is not None
assert task.category == "pcap"
# write a real sample to storage
pcap_path = os.path.join(self.binary_storage, task.sample.sha256)
shutil.copy(self.temp_pcap, pcap_path)
# reschedule the PCAP task
new_task_id = self.d.reschedule(task_id)
assert new_task_id is not None
new_task = self.d.view_task(new_task_id)
assert new_task.category == "pcap"
def test_reschedule_url(self):
# add a URL task
count = self.session.query(Task).count()
task_id = self.d.add_url("test_reschedule_url")
assert self.session.query(Task).count() == count + 1
task = self.d.view_task(task_id)
assert task is not None
assert task.category == "url"
# reschedule the URL task
new_task_id = self.d.reschedule(task_id)
assert new_task_id is not None
new_task = self.d.view_task(new_task_id)
assert new_task.category == "url"
def test_add_machine(self):
self.d.add_machine(
name="name1",
label="label",
ip="1.2.3.4",
platform="windows",
tags="tag1 tag2",
interface="int0",
snapshot="snap0",
resultserver_ip="5.6.7.8",
resultserver_port=2043,
arch="x64",
)
self.d.add_machine(
name="name2",
label="label",
ip="1.2.3.4",
platform="windows",
tags="tag1 tag2",
interface="int0",
snapshot="snap0",
resultserver_ip="5.6.7.8",
resultserver_port=2043,
arch="x64",
)
m1 = self.d.view_machine("name1")
m2 = self.d.view_machine("name2")
assert m1.to_dict() == {
"status": None,
"locked": False,
"name": "name1",
"resultserver_ip": "5.6.7.8",
"ip": "1.2.3.4",
"tags": ["tag1tag2"],
"label": "label",
"locked_changed_on": None,
"platform": "windows",
"snapshot": "snap0",
"interface": "int0",
"status_changed_on": None,
"id": 1,
"resultserver_port": "2043",
"arch": "x64",
}
assert m2.to_dict() == {
"id": 2,
"interface": "int0",
"ip": "1.2.3.4",
"label": "label",
"locked": False,
"locked_changed_on": None,
"name": "name2",
"platform": "windows",
"resultserver_ip": "5.6.7.8",
"resultserver_port": "2043",
"snapshot": "snap0",
"status": None,
"status_changed_on": None,
"tags": ["tag1tag2"],
"arch": "x64",
}