def test_set_power():
attacker = NetworkSegment(id=ns_id1, name="attacker")
target = NetworkSegment(id=ns_id2, name="target")
policy = BaseBioSegPolicy(attacker, target)
assert not policy._powers
policy.set_power(attacker, 10)
assert policy.get_power(attacker) == 10
policy.set_power(target, 20)
assert policy.get_power(target) == 20
policy.set_power(attacker, 30)
assert policy.get_power(attacker) == 30
def test_consume_objects():
profile = NetworkSegmentProfile(name="mock", is_persistent=False)
attacker = NetworkSegment(id=ns_id1, name="attacker", profile=profile)
target = NetworkSegment(id=ns_id2, name="target", profile=profile)
policy = MockBioSegPolicy(attacker, target)
assert len(policy.get_objects(attacker)) == 0
assert len(policy.get_objects(target)) == policy.N
policy.consume_objects(target, attacker)
assert len(policy.get_objects(attacker)) == policy.N
assert len(policy.get_objects(target)) == 0
assert policy.get_power(attacker) == policy.N * policy.LEVEL
assert policy.get_power(target) == 0
def test_merge_policy(a_power, a_id, t_power, t_id, expected):
"""
Test MERGE rule
:return:
"""
# Mock segments
attacker = NetworkSegment(id=a_id, name="attacker")
target = NetworkSegment(id=t_id, name="target")
policy = MergeBioSegPolicy(attacker, target)
# Mock powers
policy.set_power(attacker, a_power)
policy.set_power(target, t_power)
effective_policy = policy.get_effective_policy()
assert effective_policy
assert effective_policy.name == expected
def test_collision_floating(attacker_policy, target_policy, expected):
policy = loader.get_class(target_policy)
profile = NetworkSegmentProfile(name="stub", is_persistent=False)
target = NetworkSegment(name="target", profile=profile)
result = policy.get_effective_policy_name(target, attacker_policy)
assert result == expected
assert loader.get_class(result)
def ensure_segment(self, name):
ns = NetworkSegment.objects.filter(parent=self.object.segment.id,
name=name).first()
if not ns:
root = self.object.segment
if root.profile:
profile = root.profile.autocreated_profile or root.profile
else:
profile = None
ns = NetworkSegment(parent=self.object.segment,
name=name,
profile=profile,
description="Autocreated by segmentation")
ns.save()
return ns
def handle_split_floating(self, profile, ids, *args, **options):
connect()
p = NetworkSegmentProfile.objects.filter(name=profile).first()
if not p:
self.die("Profile not found")
if p.is_persistent:
self.die("Segment profile cannot be persistent")
for seg_id in ids:
seg = NetworkSegment.get_by_id(seg_id)
if not seg:
self.print("@@@ %s - not found. Skipping" % seg_id)
continue
self.print("@@@ Splitting %s (%s)" % (seg.name, seg_id))
objects = list(
ManagedObject.objects.filter(is_managed=True, segment=seg_id))
for mo in objects:
new_segment = NetworkSegment(
name=mo.administrative_domain.get_bioseg_floating_name(mo)
or "Bubble for %s" % mo.name,
profile=p,
parent=mo.administrative_domain.
get_bioseg_floating_parent_segment(),
)
new_segment.save()
self.print(" Moving '%s' to segment '%s'" %
(mo.name, new_segment.name))
mo.segment = new_segment
mo.save()
# Establish trials
self.print("@@@ Scheduling trials")
for mo in objects:
for link in Link.object_links(mo):
for ro in link.managed_objects:
if ro == mo:
continue
self.print(" '%s' challenging '%s' over %s -- %s" %
(mo.segment.name, ro.segment.name, mo.name,
ro.name))
BioSegTrial.schedule_trial(mo.segment,
ro.segment,
mo,
ro,
reason="link")
def get_segment(name):
ns = NetworkSegment(name=name,
parent=get_root_segment(),
profile=get_segment_profile())
ns.save()
return ns
def get_root_segment():
ns = NetworkSegment(name="Discovery Tests", profile=get_segment_profile())
ns.save()
return ns
def test_base_trial():
attacker = NetworkSegment(name="attacker")
target = NetworkSegment(name="target")
policy = BaseBioSegPolicy(attacker, target)
with pytest.raises(NotImplementedError):
policy.trial()
def test_keep_trial():
attacker = NetworkSegment(name="attacker")
target = NetworkSegment(name="target")
policy = KeepBioSegPolicy(attacker, target)
assert policy.trial() == "keep"