def test_insertBlockBetween_100perecent(render=False):
# Prep
root = Root("Start")
goal = Goal("Finish")
a = root.action("A")
b = a.action("B")
c = b.action("C")
c.connectTo(goal)
#Before
assert (len(b.edges) == 1)
# Create a second block with a pSuccess of 100
block = Block(label="Test Block", implemented=True, cost=5000, pDefend=100)
# The tested function
block.insertBetween(b, c)
# After
assert (
len(b.edges) == 1
) # As the block has 100% chance of working, c should be unreachable
assert (b.edges[0].childNode == block)
if render:
Renderer().render(root=root,
fname=inspect.currentframe().f_code.co_name)
def test_pathEvaluation(render=True):
# TODO: Make this test more useful
root = Root("Root")
goal = Goal("Goal")
brain = Brain()
# One path to vicotry (Start->a->b>c->goal)
a = root.add(Action("A", pSuccess=100, cost=500))
b = a.add(Action("B", pSuccess=50, cost=500))
c = b.add(Action("C", pSuccess=50, cost=500))
_ = c.add(goal)
# Update model with a second path to victory (Start->a->bb->goal)
bb = a.add(Action("BB", pSuccess=100, cost=20000))
bb.connectTo(goal)
paths = []
_ = Brain().pathsToVictory(root, paths)
print(paths)
assert len(paths) == 2
res = brain.evaluatePath(paths[1]) # a -> bb -> goal
assert res["attackCost"] == 20500
assert res["pSuccess"] == 100 # a * bb
if render:
Renderer().render(root=root,
fname=inspect.currentframe().f_code.co_name)
def test_instertBlockBetweenPatiallySuccessfulActions(render=True):
# Prep
root = Root("Start")
goal = Goal("Finish")
a = root.add(Action(label="A", pSuccess=80))
b = a.action("B")
c = b.action("C")
c.connectTo(goal)
# a has an 80percent chance of working.
# lets add a blocker that's 20% effective
block = Block(label="Test Block", implemented=True, cost=5000, pDefend=50)
# The tested function
block.insertBetween(a, b)
assert (len(a.edges) == 1)
assert (a.edges[0].childNode == block)
assert (len(block.edges) == 1)
assert (block.edges[0].childNode == b)
# The tested function:
if render:
Renderer().render(root=root,
fname=inspect.currentframe().f_code.co_name)
def test_insertBetween_50perecent(render=False):
# Prep
root = Root("Start")
goal = Goal("Finish")
a = root.action("A")
b = a.action("B")
c = b.action("C")
c.connectTo(goal)
# Verify prep
assert (len(a.edges) == 1)
assert (a.edges[0].childNode == b)
# Create a block
block = Block(label="Test Block", implemented=True, cost=5000, pDefend=50)
# The tested function
block.insertBetween(a, b)
# After
assert (len(a.edges) == 1)
assert (len(block.edges) == 1)
assert (a.edges[0].childNode.edges[0].childNode == b)
if render:
Renderer().render(root=root,
fname=inspect.currentframe().f_code.co_name)
def test_repr():
root = Root("root")
goal = Goal("goal")
throws = root.action("A")
dodge = throws.block("B", implemented=True)
description = throws.edges[0].__repr__()
assert description == "Edge 'Fail' connects 'A' to 'B'"
def test_detectShortcut():
root = Root("have rock")
goal = Goal("hit player")
throws = root.action("throw rock")
alarm = throws.detect("player sees rock", implemented=True)
assert isinstance(alarm, Detect)
assert len(throws.edges) == 1
edge = throws.edges[0]
assert isinstance(edge, Edge)
assert edge.parentNode == throws
assert edge.childNode == alarm
def test_single_pathToVictory(render=False):
root = Root("Root")
goal = Goal("Goal")
# Test with just one action
root.add(Action(label="A")).add(goal)
brain = Brain()
paths = brain.pathsToVictory(root)
print(paths)
assert len(paths) == 1
assert len(paths[0]) == 3
if render:
Renderer().render(root=root,
fname=inspect.currentframe().f_code.co_name)
def test_multiplePathsToVitory():
root = Root("Root")
goal = Goal("Goal")
# Test with just one action
root.add(Action(label="A")).add(goal)
assert len(root.parentEdges) == 0
assert len(goal.parentEdges) == 1
assert goal.parentEdges[0].parentNode.label == "A"
brain = Brain()
paths = brain.pathsToVictory(root)
assert len(paths) == len(goal.parentEdges)
brain = Brain()
paths = brain.pathsToVictory(root)
assert len(paths) == len(goal.parentEdges)
def test_blockShortcut():
root = Root("have rock")
goal = Goal("hit player")
throws = root.action("throw rock")
dodge = throws.block("dodge", implemented=True)
assert len(throws.edges) == 1
assert isinstance(throws.edges[0], Edge)
edge = throws.edges[0]
assert edge.childNode == dodge
assert edge.parentNode == throws
assert isinstance(dodge, Block)
assert isinstance(dodge, Node)
assert len(dodge.parentEdges) == 1
assert isinstance(dodge.parentEdges[0], Edge)
def test_double_pathToVictory(render=True):
root = Root("Root XX")
goal = Goal("Goal")
# Test with just one action
root.add(Action(label="A")).add(goal)
root.add(Action(label="B")).add(goal)
brain = Brain()
paths = brain.pathsToVictory(root)
print(paths)
assert(len(paths) == 2)
assert(len(paths[0]) == 3)
assert(len(paths[1]) == 3)
if render:
Renderer().render(
root=root,
fname=inspect.currentframe().f_code.co_name
)
def test_quad_pathToVictory(render=True):
root = Root("Root XXXX")
goal = Goal("Goal")
# Test with just one action
root.add(Action(label="A")).add(goal)
root.add(Action(label="B")).add(goal)
root.add(Action(label="C")).add(goal)
root.add(Action(label="D")).add(goal)
brain = Brain()
paths = brain.pathsToVictory(root)
assert len(paths) == 4
assert len(paths[0]) == 3
assert len(paths[1]) == 3
assert len(paths[2]) == 3
assert len(paths[3]) == 3
if render:
Renderer().render(root=root,
fname=inspect.currentframe().f_code.co_name)
def test_pathEvaluationWithBlock(render=True):
root = Root("Root")
goal = Goal("Goal")
brain = Brain()
a = root.add(Action("A",pSuccess=100, cost=500))
b = a.add(Action("b",pSuccess=70, cost=500))
b.add(goal)
paths = []
_ = brain.pathsToVictory(root, paths)
assert(len(paths)==1)
res = brain.evaluatePath(paths[0])
# Check results are correct before we add a block
assert(res['attackCost']==1000)
assert(res['pSuccess']==70)
assert(a.pSuccess==100)
block = Block(label="FIREWALL",implemented=True,cost=0,pDefend=50)
block.insertBetween(a,b)
paths = []
_ = brain.pathsToVictory(root, paths)
res = brain.evaluatePath(paths[0])
print(res)
assert(res['attackCost']==1000)
assert(res['pSuccess']==35)
if render:
Renderer().render(
root=root,
fname=inspect.currentframe().f_code.co_name
)
def basicTree():
root = Root("Root")
goal = Goal("Systems Access")
# Create three top tier actions
networkRecon = root.add(
Action(
label="Network Recon",
chain=mitreAttack["recon"],
cost=0,
time=24,
objective="Find network attack surface",
pSuccess=100,
detections=None,
))
dnsEnumeration = root.add(
Action(
label="DNS Enumeration",
chain=mitreAttack["recon"],
cost=0,
time=4,
objective="Identify all subdomains",
pSuccess=100,
))
linkedInResearch = root.add(
Action(
label="LinkedIn Research",
chain=mitreAttack["recon"],
cost=0,
time=6,
objective=
"Identify names and email addresses of current employees and customers",
pSuccess=100,
))
# Stuff learned from those activities
vpnEndpoint = networkRecon.discovery("VPN Endpoints")
sshEndpoint = networkRecon.discovery("SSH Endpoints")
subdomains = dnsEnumeration.discovery("subdomains")
employeeNames = linkedInResearch.discovery("Employee Names")
keyCustomerNames = linkedInResearch.discovery("Key Customer Details")
# Actions taken based on those discoveries
credentialStuffing = Action(
label="Credential Stuffing",
chain=mitreAttack["credStuffing"],
cost=500,
time=12,
objective="Try known username/password",
pSuccess=100,
)
_ = sshEndpoint.add(credentialStuffing)
_ = vpnEndpoint.add(credentialStuffing)
_ = employeeNames.add(credentialStuffing)
sqlmap = subdomains.add(
Action(label="sqlmap",
chain=mitreAttack["execution"],
cost=0,
time=2,
pSuccess=75))
nikto = subdomains.action("nikto")
phishing = employeeNames.action("SET Phishing")
keyCustomerNames.connectTo(phishing)
# Action Results
sqli = sqlmap.add(
Discovery(
label="Blind injection, viable RCE",
description="",
sensitivity=10,
value=1000,
))
dbExploit = sqli.add(
Action(
label="Craft & Deploy RCE",
chain=mitreAttack["execution"],
cost=0,
time=2,
pSuccess=75,
))
# Action Results that get to the goal
credentialStuffing.connectTo(goal, label="Passwords Reused")
phishing.connectTo(goal, label="Credentials Stolen")
dbExploit.connectTo(goal)
return root
# ACI Breakout Tree
# Based on the unit42 writeup here:
from attacktree.models import Action, Block, Detect, Discovery, Edge, Root, Goal
from attacktree.renderer import Renderer
root = Root("Signed up to Azure")
goal = Goal("Access to other tenants data")
whoC = root.add(Action("Deploy WhoC container to view runtime"))
omgRunC = whoC.discover("OLD runc version 1.0.0-r2")
exploit = omgRunC.action("Deploy exploit container for CVE-2019-5736")
shell = Root("Reverse shell on worker node")
exploit.add(shell)
creds = shell.discover("Read kubelet credentials from disk")
describePods = creds.action("Call KubeAPI describe pods")
pods = describePods.discover(
"100+ customer pods on 120 nodes\nEach customer has their own namespace")
versions = describePods.discover(
"OLD Kubernetes versions v1.8.4, v1.9.10, v1.10.9")
anonymousAccess = describePods.discover("Kubelets run with anonymous access")
otherKubelet = anonymousAccess.action("Access another customer's kubelet")
blockedbyFirewall = otherKubelet.block("Blocked by firewall", implemented=True)
cve = versions.discover("CVE-2018-1002102 kube-api follows 302 redirect")
exploit2 = cve.action("Attempt to redirect to kube-api pod")
fail = exploit2.block(
"ACI uses a 'bridge' POD which is not impacted by this issue",
implemented=True)
def __enter__(self):
self.root = Root(label=self.rootLabel)
self.goal = Goal(label=self.goalLabel)
return self