def test_get_by_sha256_or_create(self):
cs = ChallengeSet.create(name="foo")
assert_equals(len(IDSRule.all()), 0)
ids1 = IDSRule.get_by_sha256_or_create(rules="aaa", cs=cs)
ids2 = IDSRule.get_by_sha256_or_create(rules="aaa", cs=cs)
assert_equals(ids1, ids2)
assert_equals(len(IDSRule.all()), 1)
def test_get_or_create(self):
r = Round.create(num=0)
cs = ChallengeSet.create(name="foo")
cbn1 = ChallengeBinaryNode.create(name="foo1", cs=cs, blob="aaa1")
cbn2 = ChallengeBinaryNode.create(name="foo1", cs=cs, blob="aaa2")
ids = IDSRule.create(cs=cs, rules="aaa", sha256="sum")
cbl1, crtd1 = CSSubmissionCable.get_or_create(cs=cs,
ids=ids,
cbns=[cbn1],
round=r)
assert_true(crtd1)
cbl2, crtd2 = CSSubmissionCable.get_or_create(cs=cs,
ids=ids,
cbns=[cbn1],
round=r)
assert_false(crtd2)
assert_equals(cbl1.id, cbl2.id)
cbl3, crtd3 = CSSubmissionCable.get_or_create(cs=cs,
ids=ids,
cbns=[cbn1, cbn2],
round=r)
assert_true(crtd3)
cbl4, crtd4 = CSSubmissionCable.get_or_create(cs=cs,
ids=ids,
cbns=[cbn1, cbn2],
round=r)
assert_false(crtd4)
assert_equals(cbl3.id, cbl4.id)
def upload_cbns(args):
cs = CS.select().where(CS.name == args.cs)
patch_type, _ = PT.get_or_create(name="manual",
functionality_risk=1.0,
exploitability=0.0)
cbns = []
for patched_file in args.patched_files:
with open(patched_file) as f:
content = f.read()
try:
cbn = CBN.create(cs=cs,
blob=content,
name=patched_file,
patch_type=patch_type)
except peewee.IntegrityError:
print "CBN already exists. Fetching."
cbn = CBN.select().where(CBN.name == args.patched_file,
CBN.cs == cs, CBN.blob == content)
cbns.append(cbn)
if args.field:
ids = IDSRule.get_or_create(cs=target_cs, rules='')
CSSubmissionCable.get_or_create(cs=target_cs,
cbns=cbns,
ids=ids,
round=Round.current_round())
if args.eF is not None and args.eT is not None and args.eM is not None and cbn.patch_type is not None:
perf_score = {
'score': {
'ref': {
'task_clock': 1.0,
'rss': 1.0,
'flt': 1.0,
'file_size': 1.0
},
'rep': {
'task_clock': args.eT,
'file_size': args.size_overhead,
'rss': args.eM,
'flt': args.eM,
}
}
}
PS.create(cs=cs,
perf_score=perf_score,
patch_type=cbn.patch_type,
has_failed_polls=args.eF != 0,
failed_polls=args.eF)
if args.pT is not None and args.pM is not None and args.pS is not None:
csf.poll_feedback = PF.create(cs=cs,
round_id=Round.current_round().id,
success=args.pS,
time_overhead=args.pT,
memory_overhead=args.pM,
size_overhead=args.size_overhead)
csf.save()
def test_unsubmitted_ids_rules(self):
r1 = Round.create(num=0)
team = Team.create(name=Team.OUR_NAME)
cs = ChallengeSet.create(name="foo")
ids1 = IDSRule.create(cs=cs, rules="aaa")
ids2 = IDSRule.create(cs=cs, rules="bbb")
assert_equals(len(cs.unsubmitted_ids_rules), 2)
assert_in(ids1, cs.unsubmitted_ids_rules)
assert_in(ids2, cs.unsubmitted_ids_rules)
IDSRuleFielding.create(ids_rule=ids1, team=team, submission_round=r1)
assert_equals(len(cs.unsubmitted_ids_rules), 1)
assert_not_in(ids1, cs.unsubmitted_ids_rules)
assert_in(ids2, cs.unsubmitted_ids_rules)
IDSRuleFielding.create(ids_rule=ids2, team=team, submission_round=r1)
assert_equals(len(cs.unsubmitted_ids_rules), 0)
assert_not_in(ids1, cs.unsubmitted_ids_rules)
assert_not_in(ids2, cs.unsubmitted_ids_rules)
def field_cbns(args):
cs = CS.select().where(CS.name == args.cs)
# i know i could use one query for this, but then we might get duplicate CBNs and more than we want
cbns = [
CBN.get(CBN.cs == cs, CBN.sha256 == sha) for sha in args.patched_shas
]
ids, _ = IDSRule.get_or_create(cs=cs, rules='')
CSSubmissionCable.get_or_create(cs=cs,
cbns=cbns,
ids=ids,
round=Round.current_round())
def test_unprocessed_submission_cables(self):
r = Round.create(num=0)
cs = ChallengeSet.create(name="foo")
cbn = ChallengeBinaryNode.create(name="foo1", cs=cs, blob="aaa")
ids = IDSRule.create(cs=cs, rules="aaa", blob="aaa")
cable1 = CSSubmissionCable.create(cs=cs, ids=ids, cbns=[cbn], round=r)
cable2 = CSSubmissionCable.create(cs=cs, ids=ids, cbns=[], round=r)
assert_equals(len(cs.unprocessed_submission_cables()), 2)
assert_equals(cable1, cs.unprocessed_submission_cables()[0])
assert_equals(cable2, cs.unprocessed_submission_cables()[1])
cable1.process()
assert_equals(len(cs.unprocessed_submission_cables()), 1)
def process_patch_submission(target_cs):
"""
Process a patch submission request for the provided ChallengeSet
:param target_cs: ChallengeSet for which the request needs to be processed.
"""
round_ = Round.current_round()
cbns_to_submit = CBSubmitter.patch_decision_simple(target_cs, round_)
if cbns_to_submit is not None:
if cbns_to_submit[0].ids_rule is None:
ids = IDSRule.create(cs=target_cs, rules='')
else:
ids = cbns_to_submit[0].ids_rule
CSSubmissionCable.get_or_create(cs=target_cs,
cbns=cbns_to_submit,
ids=ids,
round=round_)
else:
LOG.info("%s - leaving old CBNs in place for", target_cs.name)
def test_generate_hash_on_create_and_save_if_missing(self):
cs = ChallengeSet.create(name="foo")
ids_create = IDSRule.create(cs=cs, rules="aaa")
assert_equals(
ids_create.sha256,
"9834876dcfb05cb167a5c24953eba58c4ac89b1adf57f28f2f9d09af107ee8f0")
ids_save = IDSRule(cs=cs, rules="bbb")
ids_save.save()
assert_equals(
ids_save.sha256,
"3e744b9dc39389baf0c5a0660589b8402f3dbb49b89b3e75f2c9355852a3c677")
ids_set = IDSRule.create(cs=cs, rules="ccc", sha256="ddd")
assert_equals(ids_set.sha256, "ddd")
def _run(self, job):
input_file = job.cbn.path
patch_type = job.payload["patch_type"]
pm = PatchMaster(input_file)
patched_bin, ids_rule = pm.create_one_patch(patch_type)
if patched_bin == None:
LOG.warning("unable to generate patch")
return
else:
name = "{}_patched_{}".format(job.cbn.name, patch_type)
ids = IDSRule.get_by_sha256_or_create(rules=ids_rule, cs=job.cbn.cs)
pt = PatchType.get(name=patch_type)
ChallengeBinaryNode.create(
root=job.cbn,
cs=job.cbn.cs,
name=name,
patch_type=pt,
blob=patched_bin,
sha256=hashlib.sha256(patched_bin).hexdigest(),
ids_rule=ids,
)
def test_variable_submitter(self):
t = Team.create(name=Team.OUR_NAME)
r0 = Round.create(num=0)
# set up several CSes
cses = [CS.create(name='CS_%s' % i) for i in range(10)]
# Set up the patches
for cs in cses:
for pt in PT.select():
ids = IDSRule.create(cs=cs, rules="HAHAHA")
cbn = CBN.create(cs=cs,
name=cs.name + "_" + pt.name,
blob="XXXX",
patch_type=pt,
ids_rule=ids)
patch_names = scriba.submitters.cb.ORIG_PATCH_ORDER
try:
cur_cssc_id = CSSC.select().order_by(CSSC.id.desc()).get().id
except CSSC.DoesNotExist:
cur_cssc_id = 0
# run the scheduler
for _ in scriba.submitters.cb.ORIG_PATCH_ORDER:
for c in cses:
scriba.submitters.cb.CBSubmitter.rotator_submission(c)
# make sure they got rotated correctly
for n, cs in enumerate(cses):
cables = list(
CSSC.select().where((CSSC.cs == cs)
& (CSSC.id > cur_cssc_id)).order_by(
CSSC.id.asc()))
assert len(cables) > 0
assert all(c.cbns[0].patch_type.name == pn
for c, pn in zip(cables, (patch_names * 10)[n:]))
def _run(self, job):
for rule_path in self._rules:
rules = open(rule_path, 'r').read()
IDSRule.create(cs=job.cs, rules=rules)