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