Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Allow loading a custom threats file #68

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
izar merged 1 commit into from
Feb 14, 2020
Merged

Allow loading a custom threats file #68

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ Element

```

For the security practitioner, you may add new threats to the `threatlib/threats.json` file:
For the security practitioner, you may supply your own threats file by setting `TM.threatsFile`. It should contain entries like:

```json
{
Expand Down
97 changes: 38 additions & 59 deletions pytm/pytm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,12 @@
'''


class varString(object):
''' A descriptor that returns strings but won't allow writing '''
def __init__(self, default):
class var(object):
''' A descriptor that allows setting a value only once '''
def __init__(self, default, onSet=None):
self.default = default
self.data = WeakKeyDictionary()
self.onSet = onSet

def __get__(self, instance, owner):
# when x.d is called we get here
Expand All @@ -33,80 +34,51 @@ def __set__(self, instance, value):
# called when x.d = val
# instance = x
# value = val
if instance in self.data:
raise ValueError(
"cannot overwrite {} value with {}, already set to {}".format(
self.__class__.__name__, value, self.data[instance]
)
)
self.data[instance] = value
if self.onSet is not None:
self.onSet(instance, value)


class varString(var):
def __set__(self, instance, value):
if not isinstance(value, str):
raise ValueError("expecting a String value, got a {}".format(type(value)))
try:
self.data[instance]
except (NameError, KeyError):
self.data[instance] = value

super().__set__(instance, value)

class varBoundary(object):
def __init__(self, default):
self.default = default
self.data = WeakKeyDictionary()

def __get__(self, instance, owner):
return self.data.get(instance, self.default)

class varBoundary(var):
def __set__(self, instance, value):
if not isinstance(value, Boundary):
raise ValueError("expecting a Boundary value, got a {}".format(type(value)))
try:
self.data[instance]
except (NameError, KeyError):
self.data[instance] = value
super().__set__(instance, value)


class varBool(object):
def __init__(self, default):
self.default = default
self.data = WeakKeyDictionary()

def __get__(self, instance, owner):
return self.data.get(instance, self.default)

class varBool(var):
def __set__(self, instance, value):
if not isinstance(value, bool):
raise ValueError("expecting a boolean value, got a {}".format(type(value)))
try:
self.data[instance]
except (NameError, KeyError):
self.data[instance] = value

super().__set__(instance, value)

class varInt(object):
def __init__(self, default):
self.default = default
self.data = WeakKeyDictionary()

def __get__(self, instance, owner):
return self.data.get(instance, self.default)

class varInt(var):
def __set__(self, instance, value):
if not isinstance(value, int):
raise ValueError("expecting an integer value, got a {}".format(type(value)))
try:
self.data[instance]
except (NameError, KeyError):
self.data[instance] = value


class varElement(object):
def __init__(self, default):
self.default = default
self.data = WeakKeyDictionary()
super().__set__(instance, value)

def __get__(self, instance, owner):
return self.data.get(instance, self.default)

class varElement(var):
def __set__(self, instance, value):
if not isinstance(value, Element):
raise ValueError("expecting an Element (or inherited) value, got a {}".format(type(value)))
try:
self.data[instance]
except (NameError, KeyError):
self.data[instance] = value
raise ValueError("expecting an Element (or inherited) "
"value, got a {}".format(type(value)))
super().__set__(instance, value)


def _setColor(element):
Expand Down Expand Up @@ -192,14 +164,22 @@ class TM():
_threatsExcluded = []
_sf = None
description = varString("")
threatsFile = varString(dirname(__file__) + "/threatlib/threats.json",
onSet=lambda i, v: i._init_threats())

def __init__(self, name, **kwargs):
for key, value in kwargs.items():
setattr(self, key, value)
self.name = name
self._sf = SuperFormatter()
# load Threats
with open(dirname(__file__) + "/threatlib/threats.json", "r", encoding="utf8") as threat_file:
self._add_threats()

def _init_threats(self):
TM._BagOfThreats = []
self._add_threats()

def _add_threats(self):
with open(self.threatsFile, "r", encoding="utf8") as threat_file:
threats_json = json.load(threat_file)

for i in threats_json:
Expand Down Expand Up @@ -547,7 +527,6 @@ class Dataflow(Element):
def __init__(self, source, sink, name, **kwargs):
self.source = source
self.sink = sink
self.name = name
super().__init__(name, **kwargs)
TM._BagOfFlows.append(self)

Expand Down
26 changes: 25 additions & 1 deletion tests/test_private_func.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
sys.path.append("..")
import unittest

from pytm.pytm import _uniq_name, Boundary
from pytm.pytm import _uniq_name, Boundary, Actor, TM


class TestUniqueNames(unittest.TestCase):
Expand All @@ -14,3 +14,27 @@ def test_duplicate_boundary_names_have_different_unique_names(self):
object_2_uniq_name = _uniq_name(object_2.name, object_2.uuid)

self.assertNotEqual(object_1_uniq_name, object_2_uniq_name)


class TestAttributes(unittest.TestCase):
def test_write_once(self):
user = Actor("User")
with self.assertRaises(ValueError):
user.name = "Computer"

def test_kwargs(self):
user = Actor("User", isAdmin=True)
self.assertEqual(user.isAdmin, True)
user = Actor("User")
self.assertEqual(user.isAdmin, False)
user.isAdmin = True
self.assertEqual(user.isAdmin, True)

def test_load_threats(self):
tm = TM("TM")
self.assertNotEqual(len(TM._BagOfThreats), 0)
with self.assertRaises(FileNotFoundError):
tm.threatsFile = "threats.json"

with self.assertRaises(FileNotFoundError):
TM("TM", threatsFile="threats.json")
10 changes: 5 additions & 5 deletions tests/test_pytmfunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,10 @@

with open(os.path.abspath(os.path.join(dirname(__file__), '..')) + "/pytm/threatlib/threats.json", "r") as threat_file:
threats_json = json.load(threat_file)

class Testpytm(unittest.TestCase):

#Test for all the threats in threats.py - test Threat.apply() function
# Test for all the threats in threats.py - test Threat.apply() function

def test_INP01(self):
lambda1 = Lambda('mylambda')
process1 = Process('myprocess')
Expand Down Expand Up @@ -407,7 +407,6 @@ def test_AC07(self):
def test_INP15(self):
web = Server("Web Server")
web.protocol = 'IMAP'
web.protocol = 'SMTP'
web.sanitizesInput = False
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "INP15"))
self.assertTrue(ThreatObj.apply(web))
Expand Down Expand Up @@ -858,5 +857,6 @@ def test_AC21(self):
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "AC21"))
self.assertTrue(ThreatObj.apply(process1))


if __name__ == '__main__':
unittest.main()
unittest.main()