def test_get_or_create(self):
assert_equals(len(ChallengeSet.select()), 0)
r1 = Round.create(num=0)
cs1, _ = ChallengeSet.get_or_create(name="foo")
cs2, _ = ChallengeSet.get_or_create(name="foo")
cs3, _ = ChallengeSet.get_or_create(name="bar")
for cs in [cs1, cs2, cs3]:
cs.rounds = [r1.id]
assert_equals(len(ChallengeSet.select()), 2)
assert_equals(cs2.id, cs1.id)
assert_not_equals(cs3.id, cs1.id)
cs3.delete_instance(recursive=True)
cs2.delete_instance(recursive=True)
cs1.delete_instance(recursive=True)
r1.delete_instance(recursive=True)
ChallengeSet._meta.database.commit()
def try_drilling(name, get_crashes):
# set up the node
cs = ChallengeSet.get_or_create(name=name.split('_')[0])
cbn = ChallengeBinaryNode.get_create(name=name, cs=cs)
cbn.root = cbn
cbn.save()
Exploit.delete().where(Exploit.cbn == cbn).execute()
if not cbn.crashes or get_crashes:
# Delete the testcases
Test.delete().where(Test.cbn == cbn).execute()
Crash.delete().where(Crash.cbn == cbn).execute()
assert len(cbn.tests) == 0
afl_job = AFLJob.create(cbn=cbn,
limit_cpu=4,
limit_memory=1,
limit_time=80)
afl_thread = AFLThread(afl_job)
afl_thread.start()
for _ in range(10):
if len(cbn.tests) == 2: break
time.sleep(1)
assert len(cbn.tests) == 2
assert len(afl_thread._worker._seen) == 2
assert len(cbn.crashes) == 0
# schedule driller
for i in cbn.tests:
dj = DrillerJob.create(payload=i,
cbn=cbn,
limit_cpu=1,
limit_memory=10)
dw = worker.DrillerWorker()
dw.run(dj)
assert len(dj.tests) > 0 #pylint:disable=no-member
for _ in range(80):
if len(cbn.crashes) > 0: break
time.sleep(1)
assert len(afl_thread._worker._seen) > 2
assert len(cbn.tests) > 2
assert len(cbn.crashes) > 0
afl_thread.join()
sorted_crashes = sorted(cbn.crashes, key=lambda x: -len(str(x.blob)))
print "Chose crash of length %d" % len(sorted_crashes[0].blob)
rj = RexJob.create(payload=sorted_crashes[0],
cbn=cbn,
limit_cpu=1,
limit_memory=10)
rw = worker.RexWorker()
rw.run(rj)
assert len(cbn.exploits) == 2
