mypy fixes

type checker doesn't like a list that contains tuples with both
length 2 and length 3. so keep length constant with None values.
This commit is contained in:
William Ballenthin
2021-06-14 09:50:12 -06:00
parent addd4683ca
commit 79459d4a14

View File

@@ -131,20 +131,18 @@ def render_attack(doc, ostream: StringIO):
if attack.get("subtechnique"):
tactics[attack["tactic"]].add((attack["technique"], attack["subtechnique"], attack["id"]))
else:
tactics[attack["tactic"]].add((attack["technique"], attack["id"]))
tactics[attack["tactic"]].add((attack["technique"], attack["id"], None))
rows = []
for tactic, techniques in sorted(tactics.items()):
inner_rows = []
for spec in sorted(techniques):
if len(spec) == 2:
technique, id = spec
if spec[-1] is None:
technique, id, _ = spec
inner_rows.append("%s %s" % (rutils.bold(technique), id))
elif len(spec) == 3:
else:
technique, subtechnique, id = spec
inner_rows.append("%s::%s %s" % (rutils.bold(technique), subtechnique, id))
else:
raise RuntimeError("unexpected ATT&CK spec format")
rows.append(
(
rutils.bold(tactic.upper()),
@@ -184,20 +182,18 @@ def render_mbc(doc, ostream: StringIO):
if mbc.get("method"):
objectives[mbc["objective"]].add((mbc["behavior"], mbc["method"], mbc["id"]))
else:
objectives[mbc["objective"]].add((mbc["behavior"], mbc["id"]))
objectives[mbc["objective"]].add((mbc["behavior"], mbc["id"], None))
rows = []
for objective, behaviors in sorted(objectives.items()):
inner_rows = []
for spec in sorted(behaviors):
if len(spec) == 2:
behavior, id = spec
if spec[-1] is None:
behavior, id, _ = spec
inner_rows.append("%s [%s]" % (rutils.bold(behavior), id))
elif len(spec) == 3:
else:
behavior, method, id = spec
inner_rows.append("%s::%s [%s]" % (rutils.bold(behavior), method, id))
else:
raise RuntimeError("unexpected MBC spec format")
rows.append(
(
rutils.bold(objective.upper()),