class DatabaseEngine(object):
"""Tests database stuff."""
URI = None
def setup_class(self):
set_cwd(tempfile.mkdtemp())
def setup(self):
self.d = Database()
self.d.connect(dsn=self.URI)
def teardown(self):
# Clear all tables without dropping them
# This is done after each test to ensure a test doesn't fail because
# of data of a previous test
meta = MetaData()
meta.reflect(self.d.engine)
ses = self.d.Session()
try:
for t in reversed(meta.sorted_tables):
ses.execute(t.delete())
ses.commit()
finally:
ses.close()
def test_add_target(self):
count = self.d.Session().query(Target).count()
add_target("http://example.com", category="url")
assert self.d.Session().query(Target).count() == count + 1
def test_add_task(self):
fd, sample_path = tempfile.mkstemp()
os.write(fd, "hehe")
os.close(fd)
# Add task.
count = self.d.Session().query(Task).count()
add_task(sample_path, category="file")
assert self.d.Session().query(Task).count() == count + 1
def test_processing_get_task(self):
# First reset all existing rows so that earlier exceptions don't affect
# this unit test run.
null, session = None, self.d.Session()
session.query(Task).filter(Task.status == "completed",
Task.processing == null).update({
"processing":
"something",
})
session.commit()
t1 = add_task("http://google.com/1",
priority=1,
status="completed",
category="url")
t2 = add_task("http://google.com/2",
priority=2,
status="completed",
category="url")
t3 = add_task("http://google.com/3",
priority=1,
status="completed",
category="url")
t4 = add_task("http://google.com/4",
priority=1,
status="completed",
category="url")
t5 = add_task("http://google.com/5",
priority=3,
status="completed",
category="url")
t6 = add_task("http://google.com/6",
priority=1,
status="completed",
category="url")
t7 = add_task("http://google.com/7",
priority=1,
status="completed",
category="url")
assert self.d.processing_get_task("foo") == t5
assert self.d.processing_get_task("foo") == t2
assert self.d.processing_get_task("foo") == t1
assert self.d.processing_get_task("foo") == t3
assert self.d.processing_get_task("foo") == t4
assert self.d.processing_get_task("foo") == t6
assert self.d.processing_get_task("foo") == t7
assert self.d.processing_get_task("foo") is None
def test_error_exists(self):
task_id = add_task("http://google.com/7", category="url")
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):
add_task("http://google.com/", category="url")
self.d.add_error("A" * 1024, 1)
err = self.d.view_errors(1)
assert err and len(err[0].message) == 1024
def test_submit(self):
dirpath = tempfile.mkdtemp()
submit_id = self.d.add_submit(dirpath, "files", {
"foo": "bar",
})
submit = self.d.view_submit(submit_id)
assert submit.id == submit_id
assert submit.tmp_path == dirpath
assert submit.submit_type == "files"
assert submit.data == {
"foo": "bar",
}
def test_connect_no_create(self):
AlembicVersion.__table__.drop(self.d.engine)
self.d.connect(dsn=self.URI, create=False)
assert "alembic_version" not in self.d.engine.table_names()
self.d.connect(dsn=self.URI)
assert "alembic_version" in self.d.engine.table_names()
def test_view_submit_tasks(self):
submit_id = self.d.add_submit(None, None, None)
target_id = add_target(__file__, category="file")
t1 = add_task(custom="1", submit_id=submit_id)
t2 = add_task(custom="2", submit_id=submit_id)
submit = self.d.view_submit(submit_id)
assert submit.id == submit_id
with pytest.raises(DetachedInstanceError):
print submit.tasks
submit = self.d.view_submit(submit_id, tasks=True)
assert len(submit.tasks) == 2
tasks = sorted((task.id, task) for task in submit.tasks)
assert tasks[0][1].id == t1
assert tasks[0][1].custom == "1"
assert tasks[1][1].id == t2
assert tasks[1][1].custom == "2"
def test_task_set_options(self):
t0 = add_task(__file__, options={"foo": "bar"})
t1 = add_task(__file__, options="foo=bar")
assert self.d.view_task(t0).options == {"foo": "bar"}
assert self.d.view_task(t1).options == {"foo": "bar"}
def test_error_action(self):
task_id = add_task(__file__)
self.d.add_error("message1", task_id)
self.d.add_error("message2", task_id, "actionhere")
e1, e2 = self.d.view_errors(task_id)
assert e1.message == "message1"
assert e1.action is None
assert e2.message == "message2"
assert e2.action == "actionhere"
def test_view_tasks(self):
t1 = add_task(__file__)
t2 = add_task("http://example.com", category="url")
tasks = self.d.view_tasks([t1, t2])
assert tasks[0].to_dict() == self.d.view_task(t1).to_dict()
assert tasks[1].to_dict() == self.d.view_task(t2).to_dict()
def test_add_machine(self):
self.d.add_machine("name1", "label", "1.2.3.4", "windows", None,
"tag1 tag2", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name2", "label", "1.2.3.4", "windows", "",
"tag1 tag2", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name3", "label", "1.2.3.4", "windows", "opt1 opt2",
"tag1 tag2", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name4",
"label",
"1.2.3.4",
"windows", ["opt3", "opt4"],
"tag1 tag2",
"int0",
"snap0",
"5.6.7.8",
2043,
"virtualbox",
reserved_by=1600)
m1 = self.d.view_machine("name1")
m2 = self.d.view_machine("name2")
m3 = self.d.view_machine("name3")
m4 = self.d.view_machine("name4")
assert m1.options == []
assert m2.options == []
assert m3.options == ["opt1", "opt2"]
assert m4.options == ["opt3", "opt4"]
assert m1.manager == "virtualbox"
assert m4.reserved_by == 1600
def test_adding_task(self):
now = datetime.datetime.now()
id = add_task(__file__, "file", 0, "py", "free=yes", 3, "custom",
"owner", "machine1", "DogeOS", ["tag1"], False, False,
now, "regular", None, now)
task = self.d.view_task(id)
assert id is not None
assert task.timeout == 0
assert task.package == "py"
assert task.options == {"free": "yes"}
assert task.priority == 3
assert task.custom == "custom"
assert task.owner == "owner"
assert task.machine == "machine1"
assert task.platform == "DogeOS"
assert len(task.tags) == 1
assert task.tags[0].name == "tag1"
assert task.memory == False
assert task.enforce_timeout == False
assert task.clock == now
assert task.submit_id is None
assert task.start_on == now
assert len(task.targets) == 1
assert task.targets[0].category == "file"
assert task.targets[0].target == __file__
def test_set_machine_rcparams(self):
self.d.add_machine("name5", "label5", "1.2.3.4", "windows", None,
"tag1 tag2", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.set_machine_rcparams("label5", {
"protocol": "rdp",
"host": "127.0.0.1",
"port": 3389,
})
m = self.d.view_machine("name5")
assert m.rcparams == {
"protocol": "rdp",
"host": "127.0.0.1",
"port": "3389",
}
def test_add_target_file(self):
fd, sample_path = tempfile.mkstemp()
os.write(fd, os.urandom(64))
os.close(fd)
target = File(sample_path)
id = add_target(sample_path, "file")
db_target = self.d.find_target(id=id)
assert id is not None
assert db_target.file_size == 64
assert db_target.file_type == target.get_type()
assert db_target.md5 == target.get_md5()
assert db_target.crc32 == target.get_crc32()
assert db_target.sha1 == target.get_sha1()
assert db_target.sha256 == target.get_sha256()
assert db_target.sha512 == target.get_sha512()
assert db_target.ssdeep == target.get_ssdeep()
assert db_target.category == "file"
def test_add_target_url(self):
target = URL("http://example.com/")
id = add_target(target.url, "url")
db_target = self.d.find_target(id=id)
assert id is not None
assert db_target.md5 == target.get_md5()
assert db_target.crc32 == target.get_crc32()
assert db_target.sha1 == target.get_sha1()
assert db_target.sha256 == target.get_sha256()
assert db_target.sha512 == target.get_sha512()
assert db_target.ssdeep == target.get_ssdeep()
assert db_target.category == "url"
def test_find_target(self):
fd, sample_path = tempfile.mkstemp()
os.write(fd, os.urandom(64))
os.close(fd)
target = File(sample_path)
id = add_target(sample_path, category="file")
assert self.d.find_target(id=id).id == id
assert self.d.find_target(crc32=target.get_crc32()).id == id
assert self.d.find_target(md5=target.get_md5()).id == id
assert self.d.find_target(sha1=target.get_sha1()).id == id
assert self.d.find_target(sha256=target.get_sha256()).id == id
assert self.d.find_target(sha512=target.get_sha512()).id == id
def test_find_target_multifilter(self):
ids = []
paths = []
target = None
for x in range(2):
fd, sample_path = tempfile.mkstemp()
randbytes = os.urandom(64)
paths.append(sample_path)
os.write(fd, randbytes)
os.close(fd)
target = File(sample_path)
ids.append(add_target(sample_path, category="file"))
db_target = self.d.find_target(sha256=target.get_sha256(),
target=paths[1])
assert self.d.find_target(id=ids[0], md5=target.get_md5()) is None
assert db_target.id == ids[1]
def test_fetch_with_machine(self):
future = datetime.datetime(2200, 5, 12, 12, 12)
add_task(__file__, category="file", tags=["service"])
t2 = add_task(__file__, category="file", machine="machine1")
add_task(__file__, category="file", start_on=future)
add_task(__file__, category="file")
t = self.d.fetch(machine="machine1", service=False)
assert t.id == t2
assert t.status == "pending"
def test_fetch_service_false(self):
add_task(__file__, category="file", tags=["service"])
t2 = add_task(__file__, category="file")
t = self.d.fetch(service=False)
assert t.id == t2
assert t.status == "pending"
def test_fetch_service_true(self):
t1 = add_task(__file__, category="file", tags=["service"])
add_task(__file__, category="file", machine="machine1")
add_task(__file__)
add_task(__file__)
task = self.d.fetch()
assert task.id == t1
assert task.status == "pending"
def test_fetch_use_start_on_true(self):
future = datetime.datetime(2200, 5, 12, 12, 12)
add_task(__file__, category="file", start_on=future, priority=999)
t2 = add_task(__file__, category="file")
t = self.d.fetch(service=False)
assert t.id == t2
assert t.status == "pending"
def test_fetch_use_start_on_false(self):
future = datetime.datetime(2200, 5, 12, 12, 12)
t1 = add_task(__file__, category="file", start_on=future, priority=999)
add_task(__file__, category="file")
t = self.d.fetch(use_start_on=False, service=False)
assert t.id == t1
assert t.status == "pending"
def test_fetch_use_exclude(self):
t1 = add_task(__file__, category="file", priority=999)
t2 = add_task(__file__, category="file", priority=999)
t3 = add_task(__file__, category="file", priority=999)
t4 = add_task(__file__, category="file", priority=999)
t = self.d.fetch(service=False, exclude=[t1, t2, t3])
assert t.id == t4
assert t.status == "pending"
def test_fetch_specific_task(self):
t1 = add_task(__file__, category="file", priority=999)
t2 = add_task(__file__, category="file", priority=999)
t = self.d.fetch(task_id=t1)
assert t.id == t1
assert t.status == "pending"
def test_lock_machine(self):
t1 = add_task(__file__, category="file", tags=["app1", "office7"])
t2 = add_task(__file__, category="file", tags=["app1", "office15"])
self.d.add_machine("name1", "name1", "1.2.3.4", "windows", "",
"app1,office7", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name2", "name2", "1.2.3.4", "DogeOS", "opt1 opt2",
"office13", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name3", "name3", "1.2.3.4", "CoffeeOS",
["opt3", "opt4"], "cofOS,office7", "int0", "snap0",
"5.6.7.8", 2043, "virtualbox")
task1 = self.d.view_task(t1)
task2 = self.d.view_task(t2)
m1 = self.d.lock_machine(tags=task1.tags)
assert m1.locked
assert m1.name == "name1"
with pytest.raises(CuckooOperationalError):
self.d.lock_machine(platform="DogeOS", tags=task2.tags)
m2 = self.d.lock_machine(platform="DogeOS")
assert m2.name == "name2"
m3 = self.d.lock_machine(label="name3")
assert m3.locked
assert m3.name == "name3"
def test_list_tasks(self):
t1 = add_task(__file__,
category="file",
owner="doge",
options={"route": "vpn511"})
t2 = add_task(__file__, category="file")
add_task(__file__, category="file")
self.d.set_status(t2, "reported")
self.d.set_status(t1, "reported")
tasks = self.d.list_tasks(owner="doge", status="reported")
tasks2 = self.d.list_tasks()
tasks3 = self.d.list_tasks(status="reported")
assert tasks[0].id == t1
assert len(tasks2) == 3
assert len(tasks3) == 2
def test_list_tasks_between(self):
for x in range(5):
add_task(__file__, category="file")
tasks = self.d.list_tasks(filter_by="id",
operators="between",
values=(1, 3))
assert len(tasks) == 3
def test_list_tasks_multiple_filter(self):
ids = []
future = None
for x in range(10):
id = add_task(__file__, category="file")
ids.append(id)
future = datetime.datetime.now() + datetime.timedelta(days=id)
ses = self.d.Session()
task = ses.query(Task).get(id)
task.completed_on = future
ses.commit()
ses.close()
tasks = self.d.list_tasks(filter_by=["id", "completed_on"],
operators=[">", "<"],
values=[4, future],
order_by="id",
limit=1)
assert len(tasks) == 1
assert tasks[0].id == 5
def test_list_tasks_offset_limit(self):
for x in range(10):
add_task(__file__, category="file")
tasks = self.d.list_tasks(offset=5, limit=10, order_by="id")
assert len(tasks) == 5
assert tasks[4].id == 10
def test_list_tasks_notvalue(self):
for x in range(10):
id = add_task(__file__, category="file")
if id % 2 == 0:
self.d.set_status(id, "running")
tasks = self.d.list_tasks(filter_by="status",
operators="!=",
values="running",
order_by="id")
assert len(tasks) == 5
assert tasks[4].id == 9
def test_list_tasks_noresults(self):
for x in range(5):
add_task(__file__, category="file")
tasks = self.d.list_tasks(status="reported")
assert tasks == []
def test_get_available_machines(self):
self.d.add_machine("name1", "name1", "1.2.3.4", "windows", "",
"app1,office7", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name2", "name2", "1.2.3.4", "DogeOS", "opt1 opt2",
"office13", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name3", "name3", "1.2.3.4", "CoffeeOS",
["opt3", "opt4"], "cofOS,office7", "int0", "snap0",
"5.6.7.8", 2043, "virtualbox")
self.d.machine_reserve(label="name2", task_id=1337)
self.d.lock_machine(label="name3")
available = self.d.get_available_machines()
names = [m["name"] for m in [db_m.to_dict() for db_m in available]]
assert len(available) == 2
assert "name2" in names
assert "name1" in names
def test_unlock_machine(self):
self.d.add_machine("name1", "name1", "1.2.3.4", "windows", "",
"app1,office7", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.lock_machine(label="name1")
assert self.d.view_machine(name="name1").locked
self.d.unlock_machine(label="name1")
assert not self.d.view_machine(name="name1").locked
def test_list_machines(self):
self.d.add_machine("name1", "name1", "1.2.3.4", "windows", "",
"app1,office7", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.add_machine("name2", "name2", "1.2.3.4", "DogeOS", "opt1 opt2",
"office13", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
allmachines = self.d.list_machines()
names = [m["name"] for m in [db_m.to_dict() for db_m in allmachines]]
assert len(allmachines) == 2
assert "name2" in names
assert "name1" in names
def test_machine_reserve(self):
self.d.add_machine("name1", "name1", "1.2.3.4", "windows", "",
"app1,office7", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
assert self.d.view_machine(name="name1").reserved_by is None
self.d.machine_reserve(label="name1", task_id=42)
assert self.d.view_machine(name="name1").reserved_by == 42
def test_clear_reservation(self):
self.d.add_machine("name1", "name1", "1.2.3.4", "windows", "",
"app1,office7", "int0", "snap0", "5.6.7.8", 2043,
"virtualbox")
self.d.machine_reserve(label="name1", task_id=42)
assert self.d.view_machine(name="name1").reserved_by == 42
self.d.clear_reservation(label="name1")
assert self.d.view_machine(name="name1").reserved_by is None
def test_clean_machines(self):
for x in range(6):
name = "name%s" % x
self.d.add_machine(name, name, "1.2.3.4", "windows", "",
"app1,office7", "int0", "snap0", "5.6.7.8",
2043, "virtualbox")
assert len(self.d.list_machines()) == 6
self.d.clean_machines()
assert len(self.d.list_machines()) == 0
def test_target_to_dict(self):
fd, sample_path = tempfile.mkstemp()
os.write(fd, os.urandom(64))
os.close(fd)
target = File(sample_path)
id = add_target(sample_path, category="file")
db_target = self.d.find_target(id=id)
db_target = db_target.to_dict()
assert db_target["id"] == id
assert db_target["file_size"] == 64
assert db_target["file_type"] == target.get_type()
assert db_target["md5"] == target.get_md5()
assert db_target["crc32"] == target.get_crc32()
assert db_target["sha1"] == target.get_sha1()
assert db_target["sha256"] == target.get_sha256()
assert db_target["sha512"] == target.get_sha512()
assert db_target["ssdeep"] == target.get_ssdeep()
assert db_target["category"] == "file"
assert db_target["target"] == sample_path
def test_task_multiple_targets(self):
db_targets = []
task_id = add_task()
for x in range(10):
fd, sample_path = tempfile.mkstemp()
os.write(fd, os.urandom(64))
os.close(fd)
add_target(sample_path, category="file", task_id=task_id)
task = self.d.view_task(task_id)
assert task.id == task_id
assert len(task.targets) == 10
def submit_tasks(target, options, package, custom, owner, timeout, priority,
machine, platform, memory, enforce_timeout, clock, tags,
remote, pattern, maxcount, is_unique, is_url, is_baseline,
is_shuffle, start_on):
db = Database()
data = dict(package=package or "",
timeout=timeout,
options=options,
priority=priority,
machine=machine,
platform=platform,
custom=custom,
owner=owner,
tags=tags,
memory="1" if memory else "0",
enforce_timeout="1" if enforce_timeout else "0",
clock=clock,
unique="1" if is_unique else "0",
start_on=start_on)
if is_baseline:
if remote:
log.info("Remote baseline support has not yet been implemented.")
return
task_id = submit_task.add_baseline(timeout, owner, machine, memory)
yield "Baseline", machine, task_id
return
if is_url and is_unique:
log.info("URL doesn't have --unique support yet.")
return
if is_url:
for url in target:
if not remote:
data.pop("unique", None)
task_id = submit_task.add_url(to_unicode(url), **data)
yield "URL", url, task_id
continue
data["url"] = to_unicode(url)
try:
r = requests.post("http://%s/tasks/create/url" % remote,
data=data)
yield "URL", url, r.json()["task_id"]
except Exception as e:
log.error("%s: unable to submit URL: %s", bold(red("Error")),
e)
else:
files = []
for path in target:
files.extend(enumerate_files(os.path.abspath(path), pattern))
if is_shuffle:
random.shuffle(files)
for filepath in files:
if not os.path.getsize(filepath):
log.warning("%s: sample %s (skipping file)",
bold(yellow("Empty")), filepath)
continue
if maxcount is not None:
if not maxcount:
break
maxcount -= 1
if not remote:
if is_unique:
sha256 = File(filepath).get_sha256()
if db.find_target(sha256=sha256):
log.info("File \"%s\" has already been analyzed",
filepath)
yield "File", filepath, None
continue
data.pop("unique", None)
task_id = submit_task.add_path(file_path=filepath, **data)
yield "File", filepath, task_id
continue
files = {
"file": (os.path.basename(filepath), open(filepath, "rb")),
}
try:
r = requests.post("http://%s/tasks/create/file" % remote,
data=data,
files=files)
yield "File", filepath, r.json()["task_id"]
except Exception as e:
log.error("%s: unable to submit file: %s", bold(red("Error")),
e)
continue