diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 4236000..cb3c940 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -39,3 +39,21 @@ jobs: run: poetry install --with dev - name: Run tests run: poetry run pytest + + lint: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v6 + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.13" + - name: Install Poetry + run: pip install "poetry>=2.1" + - name: Install dependencies + run: poetry install --with dev + - name: Check formatting + run: poetry run ruff format --check + - name: Lint + run: poetry run ruff check --output-format=github diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5a137ca..5c00506 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -49,6 +49,16 @@ can install pytm together with its development dependencies: Note that the `Makefile` targets drive the tools through Poetry, so reach for `pytest` and the `pytm` modules directly in an environment installed this way. +### Linting and formatting + +[Ruff](https://docs.astral.sh/ruff/) is both the linter and the formatter (the style is +black-compatible); the configuration lives in `pyproject.toml`. Before submitting a PR run + + make fmt + +to format the code and apply auto-fixable lint findings. CI enforces the read-only +equivalent, which you can reproduce locally with `make lint`. + ### Dependencies Dependency changes go into `pyproject.toml` - runtime ones under `[project.dependencies]`, diff --git a/Makefile b/Makefile index 58d8f47..3d279cf 100644 --- a/Makefile +++ b/Makefile @@ -63,4 +63,10 @@ docs: docs/pytm/index.html docs/threats.md .PHONY: fmt fmt: - poetry run black $(wildcard pytm/*.py) $(wildcard tests/*.py) $(wildcard *.py) + poetry run ruff check --fix --exit-zero + poetry run ruff format + +.PHONY: lint +lint: + poetry run ruff format --check + poetry run ruff check diff --git a/docs/sample_llm.py b/docs/sample_llm.py index 7c5cbac..4efcc46 100644 --- a/docs/sample_llm.py +++ b/docs/sample_llm.py @@ -1,10 +1,12 @@ #!/usr/bin/env python3 """Sample threat model demonstrating LLM element usage.""" -from pytm import TM, LLM, Server, Datastore, Boundary, Dataflow, Actor +from pytm import LLM, TM, Actor, Boundary, Dataflow, Datastore, Server tm = TM("Sample LLM Threat Model") -tm.description = "A web app using an LLM API for chat and a self-hosted model for classification" +tm.description = ( + "A web app using an LLM API for chat and a self-hosted model for classification" +) # Boundaries internet = Boundary("Internet") diff --git a/poetry.lock b/poetry.lock index 28648b7..a3fd441 100644 --- a/poetry.lock +++ b/poetry.lock @@ -12,72 +12,6 @@ files = [ {file = "annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89"}, ] -[[package]] -name = "black" -version = "26.3.1" -description = "The uncompromising code formatter." -optional = false -python-versions = ">=3.10" -groups = ["dev"] -files = [ - {file = "black-26.3.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:86a8b5035fce64f5dcd1b794cf8ec4d31fe458cf6ce3986a30deb434df82a1d2"}, - {file = "black-26.3.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5602bdb96d52d2d0672f24f6ffe5218795736dd34807fd0fd55ccd6bf206168b"}, - {file = "black-26.3.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c54a4a82e291a1fee5137371ab488866b7c86a3305af4026bdd4dc78642e1ac"}, - {file = "black-26.3.1-cp310-cp310-win_amd64.whl", hash = "sha256:6e131579c243c98f35bce64a7e08e87fb2d610544754675d4a0e73a070a5aa3a"}, - {file = "black-26.3.1-cp310-cp310-win_arm64.whl", hash = "sha256:5ed0ca58586c8d9a487352a96b15272b7fa55d139fc8496b519e78023a8dab0a"}, - {file = "black-26.3.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:28ef38aee69e4b12fda8dba75e21f9b4f979b490c8ac0baa7cb505369ac9e1ff"}, - {file = "black-26.3.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bf9bf162ed91a26f1adba8efda0b573bc6924ec1408a52cc6f82cb73ec2b142c"}, - {file = "black-26.3.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:474c27574d6d7037c1bc875a81d9be0a9a4f9ee95e62800dab3cfaadbf75acd5"}, - {file = "black-26.3.1-cp311-cp311-win_amd64.whl", hash = "sha256:5e9d0d86df21f2e1677cc4bd090cd0e446278bcbbe49bf3659c308c3e402843e"}, - {file = "black-26.3.1-cp311-cp311-win_arm64.whl", hash = "sha256:9a5e9f45e5d5e1c5b5c29b3bd4265dcc90e8b92cf4534520896ed77f791f4da5"}, - {file = "black-26.3.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e6f89631eb88a7302d416594a32faeee9fb8fb848290da9d0a5f2903519fc1"}, - {file = "black-26.3.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:41cd2012d35b47d589cb8a16faf8a32ef7a336f56356babd9fcf70939ad1897f"}, - {file = "black-26.3.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f76ff19ec5297dd8e66eb64deda23631e642c9393ab592826fd4bdc97a4bce7"}, - {file = "black-26.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:ddb113db38838eb9f043623ba274cfaf7d51d5b0c22ecb30afe58b1bb8322983"}, - {file = "black-26.3.1-cp312-cp312-win_arm64.whl", hash = "sha256:dfdd51fc3e64ea4f35873d1b3fb25326773d55d2329ff8449139ebaad7357efb"}, - {file = "black-26.3.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:855822d90f884905362f602880ed8b5df1b7e3ee7d0db2502d4388a954cc8c54"}, - {file = "black-26.3.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8a33d657f3276328ce00e4d37fe70361e1ec7614da5d7b6e78de5426cb56332f"}, - {file = "black-26.3.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f1cd08e99d2f9317292a311dfe578fd2a24b15dbce97792f9c4d752275c1fa56"}, - {file = "black-26.3.1-cp313-cp313-win_amd64.whl", hash = "sha256:c7e72339f841b5a237ff14f7d3880ddd0fc7f98a1199e8c4327f9a4f478c1839"}, - {file = "black-26.3.1-cp313-cp313-win_arm64.whl", hash = "sha256:afc622538b430aa4c8c853f7f63bc582b3b8030fd8c80b70fb5fa5b834e575c2"}, - {file = "black-26.3.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2d6bfaf7fd0993b420bed691f20f9492d53ce9a2bcccea4b797d34e947318a78"}, - {file = "black-26.3.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f89f2ab047c76a9c03f78d0d66ca519e389519902fa27e7a91117ef7611c0568"}, - {file = "black-26.3.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b07fc0dab849d24a80a29cfab8d8a19187d1c4685d8a5e6385a5ce323c1f015f"}, - {file = "black-26.3.1-cp314-cp314-win_amd64.whl", hash = "sha256:0126ae5b7c09957da2bdbd91a9ba1207453feada9e9fe51992848658c6c8e01c"}, - {file = "black-26.3.1-cp314-cp314-win_arm64.whl", hash = "sha256:92c0ec1f2cc149551a2b7b47efc32c866406b6891b0ee4625e95967c8f4acfb1"}, - {file = "black-26.3.1-py3-none-any.whl", hash = "sha256:2bd5aa94fc267d38bb21a70d7410a89f1a1d318841855f698746f8e7f51acd1b"}, - {file = "black-26.3.1.tar.gz", hash = "sha256:2c50f5063a9641c7eed7795014ba37b0f5fa227f3d408b968936e24bc0566b07"}, -] - -[package.dependencies] -click = ">=8.0.0" -mypy-extensions = ">=0.4.3" -packaging = ">=22.0" -pathspec = ">=1.0.0" -platformdirs = ">=2" -pytokens = ">=0.4.0,<0.5.0" - -[package.extras] -colorama = ["colorama (>=0.4.3)"] -d = ["aiohttp (>=3.10)"] -jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"] -uvloop = ["uvloop (>=0.15.2) ; sys_platform != \"win32\"", "winloop (>=0.5.0) ; sys_platform == \"win32\""] - -[[package]] -name = "click" -version = "8.1.8" -description = "Composable command line interface toolkit" -optional = false -python-versions = ">=3.7" -groups = ["dev"] -files = [ - {file = "click-8.1.8-py3-none-any.whl", hash = "sha256:63c132bbbed01578a06712a2d1f497bb62d9c1c0d329b7903a866228027263b2"}, - {file = "click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a"}, -] - -[package.dependencies] -colorama = {version = "*", markers = "platform_system == \"Windows\""} - [[package]] name = "colorama" version = "0.4.6" @@ -85,7 +19,7 @@ description = "Cross-platform colored terminal text." optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" groups = ["dev"] -markers = "sys_platform == \"win32\" or platform_system == \"Windows\"" +markers = "sys_platform == \"win32\"" files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, @@ -238,18 +172,6 @@ files = [ {file = "markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698"}, ] -[[package]] -name = "mypy-extensions" -version = "1.1.0" -description = "Type system extensions for programs checked with the mypy type checker." -optional = false -python-versions = ">=3.8" -groups = ["dev"] -files = [ - {file = "mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505"}, - {file = "mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558"}, -] - [[package]] name = "packaging" version = "25.0" @@ -262,23 +184,6 @@ files = [ {file = "packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f"}, ] -[[package]] -name = "pathspec" -version = "1.1.1" -description = "Utility library for gitignore style pattern matching of file paths." -optional = false -python-versions = ">=3.9" -groups = ["dev"] -files = [ - {file = "pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189"}, - {file = "pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a"}, -] - -[package.extras] -hyperscan = ["hyperscan (>=0.7)"] -optional = ["typing-extensions (>=4)"] -re2 = ["google-re2 (>=1.1)"] - [[package]] name = "pdoc3" version = "0.11.6" @@ -295,23 +200,6 @@ files = [ mako = "*" markdown = ">=3.0" -[[package]] -name = "platformdirs" -version = "4.4.0" -description = "A small Python package for determining appropriate platform-specific dirs, e.g. a `user data dir`." -optional = false -python-versions = ">=3.9" -groups = ["dev"] -files = [ - {file = "platformdirs-4.4.0-py3-none-any.whl", hash = "sha256:abd01743f24e5287cd7a5db3752faf1a2d65353f38ec26d98e25a6db65958c85"}, - {file = "platformdirs-4.4.0.tar.gz", hash = "sha256:ca753cf4d81dc309bc67b0ea38fd15dc97bc30ce419a7f58d13eb3bf14c4febf"}, -] - -[package.extras] -docs = ["furo (>=2024.8.6)", "proselint (>=0.14)", "sphinx (>=8.1.3)", "sphinx-autodoc-typehints (>=3)"] -test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=8.3.4)", "pytest-cov (>=6)", "pytest-mock (>=3.14)"] -type = ["mypy (>=1.14.1)"] - [[package]] name = "pluggy" version = "1.6.0" @@ -521,61 +409,6 @@ pygments = ">=2.7.2" [package.extras] dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "requests", "setuptools", "xmlschema"] -[[package]] -name = "pytokens" -version = "0.4.1" -description = "A Fast, spec compliant Python 3.14+ tokenizer that runs on older Pythons." -optional = false -python-versions = ">=3.8" -groups = ["dev"] -files = [ - {file = "pytokens-0.4.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2a44ed93ea23415c54f3face3b65ef2b844d96aeb3455b8a69b3df6beab6acc5"}, - {file = "pytokens-0.4.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:add8bf86b71a5d9fb5b89f023a80b791e04fba57960aa790cc6125f7f1d39dfe"}, - {file = "pytokens-0.4.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:670d286910b531c7b7e3c0b453fd8156f250adb140146d234a82219459b9640c"}, - {file = "pytokens-0.4.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:4e691d7f5186bd2842c14813f79f8884bb03f5995f0575272009982c5ac6c0f7"}, - {file = "pytokens-0.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:27b83ad28825978742beef057bfe406ad6ed524b2d28c252c5de7b4a6dd48fa2"}, - {file = "pytokens-0.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d70e77c55ae8380c91c0c18dea05951482e263982911fc7410b1ffd1dadd3440"}, - {file = "pytokens-0.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a58d057208cb9075c144950d789511220b07636dd2e4708d5645d24de666bdc"}, - {file = "pytokens-0.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b49750419d300e2b5a3813cf229d4e5a4c728dae470bcc89867a9ad6f25a722d"}, - {file = "pytokens-0.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d9907d61f15bf7261d7e775bd5d7ee4d2930e04424bab1972591918497623a16"}, - {file = "pytokens-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:ee44d0f85b803321710f9239f335aafe16553b39106384cef8e6de40cb4ef2f6"}, - {file = "pytokens-0.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:140709331e846b728475786df8aeb27d24f48cbcf7bcd449f8de75cae7a45083"}, - {file = "pytokens-0.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d6c4268598f762bc8e91f5dbf2ab2f61f7b95bdc07953b602db879b3c8c18e1"}, - {file = "pytokens-0.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:24afde1f53d95348b5a0eb19488661147285ca4dd7ed752bbc3e1c6242a304d1"}, - {file = "pytokens-0.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5ad948d085ed6c16413eb5fec6b3e02fa00dc29a2534f088d3302c47eb59adf9"}, - {file = "pytokens-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:3f901fe783e06e48e8cbdc82d631fca8f118333798193e026a50ce1b3757ea68"}, - {file = "pytokens-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8bdb9d0ce90cbf99c525e75a2fa415144fd570a1ba987380190e8b786bc6ef9b"}, - {file = "pytokens-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5502408cab1cb18e128570f8d598981c68a50d0cbd7c61312a90507cd3a1276f"}, - {file = "pytokens-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:29d1d8fb1030af4d231789959f21821ab6325e463f0503a61d204343c9b355d1"}, - {file = "pytokens-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:970b08dd6b86058b6dc07efe9e98414f5102974716232d10f32ff39701e841c4"}, - {file = "pytokens-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:9bd7d7f544d362576be74f9d5901a22f317efc20046efe2034dced238cbbfe78"}, - {file = "pytokens-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4a14d5f5fc78ce85e426aa159489e2d5961acf0e47575e08f35584009178e321"}, - {file = "pytokens-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97f50fd18543be72da51dd505e2ed20d2228c74e0464e4262e4899797803d7fa"}, - {file = "pytokens-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc74c035f9bfca0255c1af77ddd2d6ae8419012805453e4b0e7513e17904545d"}, - {file = "pytokens-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f66a6bbe741bd431f6d741e617e0f39ec7257ca1f89089593479347cc4d13324"}, - {file = "pytokens-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:b35d7e5ad269804f6697727702da3c517bb8a5228afa450ab0fa787732055fc9"}, - {file = "pytokens-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8fcb9ba3709ff77e77f1c7022ff11d13553f3c30299a9fe246a166903e9091eb"}, - {file = "pytokens-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:79fc6b8699564e1f9b521582c35435f1bd32dd06822322ec44afdeba666d8cb3"}, - {file = "pytokens-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d31b97b3de0f61571a124a00ffe9a81fb9939146c122c11060725bd5aea79975"}, - {file = "pytokens-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:967cf6e3fd4adf7de8fc73cd3043754ae79c36475c1c11d514fc72cf5490094a"}, - {file = "pytokens-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918"}, - {file = "pytokens-0.4.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:da5baeaf7116dced9c6bb76dc31ba04a2dc3695f3d9f74741d7910122b456edc"}, - {file = "pytokens-0.4.1-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:11edda0942da80ff58c4408407616a310adecae1ddd22eef8c692fe266fa5009"}, - {file = "pytokens-0.4.1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0fc71786e629cef478cbf29d7ea1923299181d0699dbe7c3c0f4a583811d9fc1"}, - {file = "pytokens-0.4.1-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:dcafc12c30dbaf1e2af0490978352e0c4041a7cde31f4f81435c2a5e8b9cabb6"}, - {file = "pytokens-0.4.1-cp38-cp38-win_amd64.whl", hash = "sha256:42f144f3aafa5d92bad964d471a581651e28b24434d184871bd02e3a0d956037"}, - {file = "pytokens-0.4.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:34bcc734bd2f2d5fe3b34e7b3c0116bfb2397f2d9666139988e7a3eb5f7400e3"}, - {file = "pytokens-0.4.1-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:941d4343bf27b605e9213b26bfa1c4bf197c9c599a9627eb7305b0defcfe40c1"}, - {file = "pytokens-0.4.1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3ad72b851e781478366288743198101e5eb34a414f1d5627cdd585ca3b25f1db"}, - {file = "pytokens-0.4.1-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:682fa37ff4d8e95f7df6fe6fe6a431e8ed8e788023c6bcc0f0880a12eab80ad1"}, - {file = "pytokens-0.4.1-cp39-cp39-win_amd64.whl", hash = "sha256:30f51edd9bb7f85c748979384165601d028b84f7bd13fe14d3e065304093916a"}, - {file = "pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de"}, - {file = "pytokens-0.4.1.tar.gz", hash = "sha256:292052fe80923aae2260c073f822ceba21f3872ced9a68bb7953b348e561179a"}, -] - -[package.extras] -dev = ["black", "build", "mypy", "pytest", "pytest-cov", "setuptools", "tox", "twine", "wheel"] - [[package]] name = "ruff" version = "0.15.11" @@ -634,4 +467,4 @@ typing-extensions = ">=4.12.0" [metadata] lock-version = "2.1" python-versions = ">=3.11,<3.15" -content-hash = "f4238c8a1f5b27910133f8160430adf18273c52749fbd59c9250c3909fb35111" +content-hash = "0e0c028892b6b547fb64ef4cdd048b6db4a444251fa2355c7199db9501f807b4" diff --git a/pyproject.toml b/pyproject.toml index 83aadc0..5dfea9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,6 @@ Homepage = "https://github.com/OWASP/pytm" [dependency-groups] dev = [ "pytest>=8.3.5,<10.0.0", - "black>=25.9,<27.0", "pdoc3>=0.11.6,<0.12.0", "ruff>=0.15.11,<0.16.0", ] @@ -42,8 +41,7 @@ select = [ "UP", # pyupgrade ] ignore = [ - "E501", # line too long — black handles line length for reformattable code - "UP007", # Union[X, Y] -> X | Y syntax requires Python 3.10+ at runtime + "E501", # line too long — the formatter handles line length for reformattable code ] [build-system] diff --git a/pytm/__init__.py b/pytm/__init__.py index 6e1ee95..ee5ac29 100644 --- a/pytm/__init__.py +++ b/pytm/__init__.py @@ -29,22 +29,22 @@ import sys -from .json import load, loads -from .pytm import var +from .actor import Actor +from .asset import LLM, Agent, Asset, ExternalEntity, Lambda, Server +from .base import Assumption, Controls +from .boundary import Boundary +from .data import Data +from .dataflow import Dataflow +from .datastore import Datastore +from .element import Element # Import from new Pydantic models from .enums import Action, Classification, DatastoreType, Lifetime, TLSVersion -from .base import Assumption, Controls -from .element import Element -from .data import Data -from .threat import Threat from .finding import Finding -from .asset import Agent, Asset, Lambda, LLM, Server, ExternalEntity -from .datastore import Datastore -from .actor import Actor +from .json import load, loads from .process import Process, SetOfProcesses -from .dataflow import Dataflow -from .boundary import Boundary +from .pytm import var +from .threat import Threat from .tm import TM # Rebuild models to resolve forward references diff --git a/pytm/actor.py b/pytm/actor.py index c46c72e..bbde64a 100644 --- a/pytm/actor.py +++ b/pytm/actor.py @@ -1,10 +1,11 @@ """Actor model - represents entities that initiate actions.""" -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING + from pydantic import Field, field_validator -from .element import Element from .base import DataSet +from .element import Element if TYPE_CHECKING: from .dataflow import Dataflow @@ -35,10 +36,10 @@ class Actor(Element): default_factory=DataSet, description="pytm.Data object(s) in outgoing data flows", ) - inputs: List["Dataflow"] = Field( + inputs: list["Dataflow"] = Field( default_factory=list, description="Incoming Dataflows" ) - outputs: List["Dataflow"] = Field( + outputs: list["Dataflow"] = Field( default_factory=list, description="Outgoing Dataflows" ) isAdmin: bool = Field( diff --git a/pytm/asset.py b/pytm/asset.py index 33a9b61..296634b 100644 --- a/pytm/asset.py +++ b/pytm/asset.py @@ -1,11 +1,11 @@ """Asset models - base Asset class and specific asset implementations.""" -from typing import List, TYPE_CHECKING +from typing import TYPE_CHECKING from pydantic import Field, field_validator -from .element import Element, sev_to_color from .base import DataSet +from .element import Element, sev_to_color if TYPE_CHECKING: from .dataflow import Dataflow @@ -36,10 +36,10 @@ class Asset(Element): default_factory=DataSet, description="pytm.Data object(s) in incoming data flows", ) - inputs: List["Dataflow"] = Field( + inputs: list["Dataflow"] = Field( default_factory=list, description="incoming Dataflows" ) - outputs: List["Dataflow"] = Field( + outputs: list["Dataflow"] = Field( default_factory=list, description="outgoing Dataflows" ) onAWS: bool = Field(default=False, description="Is this asset on AWS?") diff --git a/pytm/base.py b/pytm/base.py index d72214d..badac69 100644 --- a/pytm/base.py +++ b/pytm/base.py @@ -2,15 +2,16 @@ from __future__ import annotations -from typing import Any, Iterable, List, Set, Union, TYPE_CHECKING +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any from pydantic import BaseModel, ConfigDict, Field if TYPE_CHECKING: - from .element import Element from .data import Data - from .threat import Threat + from .element import Element from .finding import Finding + from .threat import Threat class DataSet(set): @@ -18,9 +19,9 @@ class DataSet(set): __slots__ = ("_names",) - def __init__(self, values: Iterable["Data"] | None = None): + def __init__(self, values: Iterable[Data] | None = None): super().__init__() - self._names: Set[str] = set() + self._names: set[str] = set() if values is not None: self.update(values) @@ -74,12 +75,12 @@ def clear(self) -> None: # type: ignore[override] super().clear() self._names.clear() - def _register(self, element: "Data") -> None: + def _register(self, element: Data) -> None: name = getattr(element, "name", None) if isinstance(name, str): self._names.add(name) - def _unregister(self, element: "Data") -> None: + def _unregister(self, element: Data) -> None: name = getattr(element, "name", None) if isinstance(name, str): self._names.discard(name) @@ -177,7 +178,7 @@ class Assumption(BaseModel): model_config = ConfigDict(extra="allow") name: str = Field(description="Name of the assumption") - exclude: Set[str] = Field( + exclude: set[str] = Field( default_factory=set, description="A set of threat SIDs to exclude for this assumption. For example: INP01", ) @@ -186,7 +187,7 @@ class Assumption(BaseModel): ) def __init__( - self, name: str = None, exclude: Union[List[str], Set[str]] = None, **kwargs + self, name: str = None, exclude: list[str] | set[str] = None, **kwargs ): """Initialize an Assumption. @@ -208,9 +209,9 @@ def __str__(self): # Type aliases for complex field types that reference forward declarations -ElementList = List["Element"] -DataList = List["Data"] -ThreatList = List["Threat"] -FindingList = List["Finding"] +ElementList = list["Element"] +DataList = list["Data"] +ThreatList = list["Threat"] +FindingList = list["Finding"] ControlsType = Controls -AssumptionList = List[Assumption] +AssumptionList = list[Assumption] diff --git a/pytm/boundary.py b/pytm/boundary.py index 5a87006..d5adc68 100644 --- a/pytm/boundary.py +++ b/pytm/boundary.py @@ -1,7 +1,7 @@ """Boundary model - represents trust boundaries in the threat model.""" -from typing import List, TYPE_CHECKING from textwrap import indent +from typing import TYPE_CHECKING from .element import Element @@ -75,7 +75,7 @@ def _color(self, **kwargs) -> str: else: return "firebrick2" - def parents(self) -> List["Boundary"]: + def parents(self) -> list["Boundary"]: """Get parent boundaries.""" result = [] parent = self.inBoundary diff --git a/pytm/data.py b/pytm/data.py index c11f24a..3b7a0fa 100644 --- a/pytm/data.py +++ b/pytm/data.py @@ -1,13 +1,14 @@ """Data model - represents data that traverses the threat model.""" -from typing import List, TYPE_CHECKING -from pydantic import BaseModel, Field, ConfigDict +from typing import TYPE_CHECKING + +from pydantic import BaseModel, ConfigDict, Field from .enums import Classification, Lifetime if TYPE_CHECKING: - from .element import Element from .dataflow import Dataflow + from .element import Element class Data(BaseModel): @@ -59,10 +60,10 @@ class Data(BaseModel): isSourceEncryptedAtRest: bool = Field( default=False, description="Is data encrypted at rest at source?" ) - carriedBy: List["Dataflow"] = Field( + carriedBy: list["Dataflow"] = Field( default_factory=list, description="Dataflows that carries this piece of data" ) - processedBy: List["Element"] = Field( + processedBy: list["Element"] = Field( default_factory=list, description="Elements that store/process this piece of data", ) diff --git a/pytm/dataflow.py b/pytm/dataflow.py index 3d7315a..c1f6d70 100644 --- a/pytm/dataflow.py +++ b/pytm/dataflow.py @@ -1,11 +1,12 @@ """Dataflow model - represents data flows between elements.""" from typing import Optional + from pydantic import Field, field_validator, model_validator +from .base import DataSet from .element import Element, sev_to_color from .enums import Classification, TLSVersion -from .base import DataSet class Dataflow(Element): diff --git a/pytm/datastore.py b/pytm/datastore.py index 0493406..852e4bb 100644 --- a/pytm/datastore.py +++ b/pytm/datastore.py @@ -2,6 +2,7 @@ import os from typing import TYPE_CHECKING + from pydantic import Field from .asset import Asset diff --git a/pytm/element.py b/pytm/element.py index 6212fdb..38f9a60 100644 --- a/pytm/element.py +++ b/pytm/element.py @@ -5,7 +5,7 @@ import uuid as uuid_module from hashlib import sha224 from textwrap import wrap -from typing import Any, List, Optional, Set, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Optional from pydantic import BaseModel, ConfigDict, Field, field_validator @@ -72,23 +72,23 @@ class Element(BaseModel): default=TLSVersion.NONE, description="Minimum TLS version required", ) - findings: List["Finding"] = Field( + findings: list["Finding"] = Field( default_factory=list, description="Threats that apply to this element", ) - overrides: List["Finding"] = Field( + overrides: list["Finding"] = Field( default_factory=list, description="Overrides to findings, allowing to set a custom response, CVSS score or override other attributes", ) - assumptions: List[Assumption] = Field( + assumptions: list[Assumption] = Field( default_factory=list, description="Assumptions about the element. These optionally allow to exclude threats with the given SIDs", ) - levels: Set[int] = Field( + levels: set[int] = Field( default_factory=lambda: {0}, description="List of levels (0, 1, 2, ...) to be drawn in the model", ) - sourceFiles: List[str] = Field( + sourceFiles: list[str] = Field( default_factory=list, description="Location of the source code that describes this element relative to the directory of the model script", ) @@ -119,9 +119,7 @@ def _coerce_levels(cls, value): return set(value) return {value} - def __setattr__( - self, key: str, value: Any - ) -> None: # noqa: D401 - keep same behaviour + def __setattr__(self, key: str, value: Any) -> None: # noqa: D401 - keep same behaviour if ( key in self._WRITE_ONCE_FIELDS and key in self.__dict__ @@ -130,7 +128,7 @@ def __setattr__( raise ValueError(f"cannot overwrite {type(self).__name__}.{key} value") super().__setattr__(key, value) - def __init__(self, name: Optional[str] = None, **data: Any): + def __init__(self, name: str | None = None, **data: Any): """Initialize an Element. Args: @@ -300,7 +298,7 @@ def _attr_values(self) -> dict: """Return a dictionary of all attribute values.""" return self.model_dump() - def checkTLSVersion(self, flows: List["Dataflow"]) -> bool: + def checkTLSVersion(self, flows: list["Dataflow"]) -> bool: """Check if any flows have insufficient TLS version.""" return any(f.tlsVersion < self.minTLSVersion for f in flows) diff --git a/pytm/finding.py b/pytm/finding.py index 92bd8e5..e1e392a 100644 --- a/pytm/finding.py +++ b/pytm/finding.py @@ -1,7 +1,8 @@ """Finding model - represents a finding linking an element to a threat.""" -from typing import Optional, TYPE_CHECKING -from pydantic import BaseModel, Field, ConfigDict +from typing import TYPE_CHECKING, Optional + +from pydantic import BaseModel, ConfigDict, Field from .base import Assumption @@ -49,7 +50,7 @@ class Finding(BaseModel): threat_id: str = Field(description="Threat ID") references: str = Field(description="Threat references") condition: str = Field(description="Threat condition") - assumption: Optional[Assumption] = Field( + assumption: Assumption | None = Field( default=None, description="The assumption that caused this finding to be excluded", ) diff --git a/pytm/flows.py b/pytm/flows.py index 9ef3d85..b1ae551 100644 --- a/pytm/flows.py +++ b/pytm/flows.py @@ -2,7 +2,9 @@ from pytm import Element -def req_reply(src: Element, dest: Element, req_name: str, reply_name=None) -> tuple[DF, DF]: +def req_reply( + src: Element, dest: Element, req_name: str, reply_name=None +) -> tuple[DF, DF]: """ This function creates two datflows where one dataflow is a request and the second dataflow is the corresponding reply to the newly created request. diff --git a/pytm/json.py b/pytm/json.py index be54d7d..129e3df 100644 --- a/pytm/json.py +++ b/pytm/json.py @@ -1,15 +1,15 @@ import json -from .tm import TM +from .actor import Actor +from .asset import LLM, Agent, Asset, ExternalEntity, Lambda, Server from .base import DataSet from .boundary import Boundary from .data import Data from .dataflow import Dataflow -from .asset import Agent, Asset, Server, ExternalEntity, Lambda, LLM from .datastore import Datastore -from .actor import Actor -from .process import Process, SetOfProcesses from .enums import Action, Classification, Lifetime +from .process import Process, SetOfProcesses +from .tm import TM _ELEMENT_CLASSES = { "Asset": Asset, diff --git a/pytm/process.py b/pytm/process.py index e71cb7a..268f3e3 100644 --- a/pytm/process.py +++ b/pytm/process.py @@ -1,6 +1,7 @@ """Process model - represents processes that handle data.""" from typing import TYPE_CHECKING + from pydantic import Field from .asset import Asset diff --git a/pytm/pytm.py b/pytm/pytm.py index 9a5fb7c..29e9058 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -1,19 +1,26 @@ import argparse -import html import copy +import html import logging import re import sys - +from collections import defaultdict +from collections.abc import Iterable, Mapping from dataclasses import dataclass, field +from functools import singledispatch from typing import ClassVar from pydantic import ValidationError from pydantic_core import PydanticUndefined -from collections import defaultdict -from collections.abc import Iterable, Mapping -from functools import singledispatch +from .actor import Actor +from .asset import LLM, Agent, Asset, ExternalEntity, Lambda, Server +from .base import Assumption, Controls +from .boundary import Boundary +from .data import Data +from .dataflow import Dataflow +from .datastore import Datastore +from .element import Element # Import all the new Pydantic models from .enums import ( @@ -21,20 +28,12 @@ Classification, DatastoreType, Lifetime, - TLSVersion, OrderedEnum, + TLSVersion, ) -from .base import Assumption, Controls -from .element import Element -from .data import Data -from .threat import Threat from .finding import Finding -from .asset import Agent, Asset, Lambda, LLM, Server, ExternalEntity -from .datastore import Datastore -from .actor import Actor from .process import Process, SetOfProcesses -from .dataflow import Dataflow -from .boundary import Boundary +from .threat import Threat from .tm import TM, UIError logger = logging.getLogger(__name__) @@ -184,7 +183,7 @@ def _describe_classes(class_names): print(" (no attributes)") else: longest = len(max(field_names, key=len)) + 2 - lpadding = f'\n{" ":<{longest+2}}' + lpadding = f"\n{' ':<{longest + 2}}" for field_name in field_names: field_info = model_fields[field_name] docs: list[str] = [] @@ -222,7 +221,7 @@ def _describe_classes(class_names): print(" (no attributes)") else: longest = len(max(attrs, key=len)) + 2 - lpadding = f'\n{" ":<{longest+2}}' + lpadding = f"\n{' ':<{longest + 2}}" for attr in sorted(attrs): value = getattr(klass, attr) docs = [] diff --git a/pytm/report_util.py b/pytm/report_util.py index e0994de..4d52b14 100644 --- a/pytm/report_util.py +++ b/pytm/report_util.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any, List +from typing import Any class ReportUtils: @@ -22,7 +22,7 @@ def getParentName(element: Any) -> str: return parent.name if parent is not None else "" @staticmethod - def getNamesOfParents(element: Any) -> List[str] | str: + def getNamesOfParents(element: Any) -> list[str] | str: """Return a list of parent boundary names for *element*.""" from pytm import Boundary diff --git a/pytm/template_engine.py b/pytm/template_engine.py index da2c6d6..72380dc 100644 --- a/pytm/template_engine.py +++ b/pytm/template_engine.py @@ -5,17 +5,15 @@ from __future__ import annotations import string -from collections.abc import Iterable -from functools import lru_cache -from typing import Any, Callable +from collections.abc import Callable, Iterable +from functools import cache +from typing import Any class SuperFormatter(string.Formatter): """Lightweight formatter with helpers for reports and templates.""" - def format_field( - self, value: Any, format_spec: str - ) -> Any: # noqa: D401 - same semantics as base + def format_field(self, value: Any, format_spec: str) -> Any: # noqa: D401 - same semantics as base if not format_spec: return super().format_field(value, format_spec) @@ -70,7 +68,7 @@ def call_util_method(self, method_name: str, obj: Any) -> Any: return method(obj) @staticmethod - @lru_cache(maxsize=None) + @cache def _resolve_report_method(method_name: str) -> Callable[[Any], Any]: from pytm.report_util import ReportUtils diff --git a/pytm/threat.py b/pytm/threat.py index f5f13c7..0107289 100644 --- a/pytm/threat.py +++ b/pytm/threat.py @@ -3,19 +3,18 @@ from __future__ import annotations import ast +import builtins import sys -from types import CodeType -from typing import Any, ClassVar, Tuple, List from collections.abc import Iterable - -import builtins +from types import CodeType +from typing import Any, ClassVar from pydantic import ( BaseModel, - Field, ConfigDict, - model_validator, + Field, PrivateAttr, + model_validator, ) @@ -130,8 +129,8 @@ def visit_Name(self, node: ast.Name) -> Any: # noqa: D401 return None @staticmethod - def _attribute_chain(node: ast.Attribute) -> List[str]: - chain: List[str] = [node.attr] + def _attribute_chain(node: ast.Attribute) -> list[str]: + chain: list[str] = [node.attr] current = node.value while isinstance(current, ast.Attribute): if isinstance(current.attr, str) and current.attr.startswith("__"): @@ -190,7 +189,7 @@ class Threat(BaseModel): prerequisites: str = Field(default="", description="Prerequisites for the threat") example: str = Field(default="", description="Example of the threat") references: str = Field(default="", description="References for the threat") - target: Tuple = Field(default=(), description="Target classes for this threat") + target: tuple = Field(default=(), description="Target classes for this threat") _compiled_condition: CodeType | None = PrivateAttr(default=None) _eval_globals: ClassVar[dict[str, Any] | None] = None diff --git a/pytm/tm.py b/pytm/tm.py index 891dee5..40ca0d2 100644 --- a/pytm/tm.py +++ b/pytm/tm.py @@ -9,31 +9,32 @@ import random import re import sys -from collections import defaultdict, Counter +from collections import Counter, defaultdict +from collections.abc import Iterable from dataclasses import dataclass, field from datetime import datetime +from html import escape as html_escape from itertools import combinations from textwrap import indent -from typing import ClassVar, Dict, Iterable, List, TYPE_CHECKING -from html import escape as html_escape +from typing import TYPE_CHECKING, ClassVar -from pydantic import BaseModel, Field, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator -from .enums import Action from .base import Assumption +from .enums import Action from .template_engine import SuperFormatter logger = logging.getLogger(__name__) if TYPE_CHECKING: - from .element import Element - from .asset import Asset from .actor import Actor - from .dataflow import Dataflow + from .asset import Asset from .boundary import Boundary from .data import Data - from .threat import Threat + from .dataflow import Dataflow + from .element import Element from .finding import Finding + from .threat import Threat class UIError(Exception): @@ -48,14 +49,14 @@ def __init__(self, e, context): class TMState: """Mutable registry for TM-owned collections.""" - flows: List["Dataflow"] = field(default_factory=list) - elements: List["Element"] = field(default_factory=list) - actors: List["Actor"] = field(default_factory=list) - assets: List["Asset"] = field(default_factory=list) - threats: List["Threat"] = field(default_factory=list) - boundaries: List["Boundary"] = field(default_factory=list) - data: List["Data"] = field(default_factory=list) - threats_excluded: List[str] = field(default_factory=list) + flows: list[Dataflow] = field(default_factory=list) + elements: list[Element] = field(default_factory=list) + actors: list[Actor] = field(default_factory=list) + assets: list[Asset] = field(default_factory=list) + threats: list[Threat] = field(default_factory=list) + boundaries: list[Boundary] = field(default_factory=list) + data: list[Data] = field(default_factory=list) + threats_excluded: list[str] = field(default_factory=list) class _StateAttribute: @@ -63,7 +64,7 @@ class _StateAttribute: def __init__(self, field_name: str): self.field_name = field_name - self.owner: type["TM"] | None = None + self.owner: type[TM] | None = None def __set_name__(self, owner, name): self.owner = owner @@ -104,7 +105,7 @@ class TM(BaseModel, metaclass=TMModelMetaclass): ) _state: ClassVar[TMState] = TMState() - _state_attributes: ClassVar[Dict[str, _StateAttribute]] = {} + _state_attributes: ClassVar[dict[str, _StateAttribute]] = {} _duplicate_ignored_attrs: ClassVar[tuple[str, ...]] = () @classmethod @@ -140,10 +141,10 @@ def _get_state(cls) -> TMState: ignoreUnused: bool = Field( default=False, description="Ignore elements not used in any Dataflow" ) - findings: List["Finding"] = Field( + findings: list[Finding] = Field( default_factory=list, description="Threats found for elements of this model" ) - excluded_findings: List["Finding"] = Field( + excluded_findings: list[Finding] = Field( default_factory=list, description="Threats found for elements of this model, that were excluded on a per-element basis, using the Assumptions class", ) @@ -151,7 +152,7 @@ def _get_state(cls) -> TMState: default=Action.NO_ACTION, description="How to handle duplicate Dataflow with same properties, except name and notes", ) - assumptions: List[Assumption] = Field( + assumptions: list[Assumption] = Field( default_factory=list, description="A list of assumptions about the design/model" ) colormap: bool = Field(default=False, exclude=True) @@ -242,7 +243,7 @@ def _init_threats(self): def _add_threats(self): """Add threats from the threats file.""" try: - with open(self.threatsFile, "r", encoding="utf8") as threat_file: + with open(self.threatsFile, encoding="utf8") as threat_file: threats_json = json.load(threat_file) except (FileNotFoundError, PermissionError, IsADirectoryError) as e: raise UIError( @@ -305,9 +306,10 @@ def check(self): def resolve(self): """Resolve threats and generate findings.""" - from .finding import Finding from collections import defaultdict + from .finding import Finding + finding_count = 0 excluded_finding_count = 0 findings = [] @@ -380,7 +382,7 @@ def process(self): try: self._process() except UIError as e: # pragma: no cover - mirrors historical behaviour - message = "Failed to execute\n" f" {e.context}\n" f" {e.error}\n" + message = f"Failed to execute\n {e.context}\n {e.error}\n" sys.stderr.write(message) raise SystemExit(127) from e @@ -522,6 +524,7 @@ def _dfd_template(self): def dfd(self, **kwargs): """Generate Data Flow Diagram.""" from collections import defaultdict + from .boundary import Boundary if "levels" in kwargs: @@ -582,42 +585,42 @@ def _seq_template(self): def seq(self): """Generate sequence diagram.""" from .actor import Actor - from .datastore import Datastore from .boundary import Boundary from .dataflow import Dataflow + from .datastore import Datastore participants = [] for e in TM._elements: if isinstance(e, Actor): participants.append( - 'actor {0} as "{1}"'.format( + 'actor {} as "{}"'.format( e._uniq_name(), getattr(e, "display_name", lambda: e.name)() ) ) elif isinstance(e, Datastore): participants.append( - 'database {0} as "{1}"'.format( + 'database {} as "{}"'.format( e._uniq_name(), getattr(e, "display_name", lambda: e.name)() ) ) elif not isinstance(e, (Dataflow, Boundary)): participants.append( - 'entity {0} as "{1}"'.format( + 'entity {} as "{}"'.format( e._uniq_name(), getattr(e, "display_name", lambda: e.name)() ) ) messages = [] for e in TM._flows: - message = "{0} -> {1}: {2}".format( + message = "{} -> {}: {}".format( e.source._uniq_name(), e.sink._uniq_name(), getattr(e, "display_name", lambda: e.name)(), ) note = "" if getattr(e, "note", "") != "": - note = "\nnote left\n{}\nend note".format(e.note) - messages.append("{}{}".format(message, note)) + note = f"\nnote left\n{e.note}\nend note" + messages.append(f"{message}{note}") return self._seq_template().format( participants="\n".join(participants), messages="\n".join(messages) @@ -792,12 +795,7 @@ def _check_duplicates(self, flows): right.is_drawn = True continue raise ValueError( - "Duplicate Dataflow found between {} and {}: {} is same as {}".format( - left.source, - left.sink, - left, - right, - ) + f"Duplicate Dataflow found between {left.source} and {left.sink}: {left} is same as {right}" ) diff --git a/tests/test_flows_helpers.py b/tests/test_flows_helpers.py index dc5b16e..bff9779 100644 --- a/tests/test_flows_helpers.py +++ b/tests/test_flows_helpers.py @@ -1,6 +1,6 @@ import pytest -from pytm import Process, Server, TM +from pytm import TM, Process, Server from pytm.dataflow import Dataflow from pytm.flows import reply, req_reply diff --git a/tests/test_private_func.py b/tests/test_private_func.py index 7f7b06e..95e4eed 100644 --- a/tests/test_private_func.py +++ b/tests/test_private_func.py @@ -1,4 +1,5 @@ import random + import pytest from pytm.pytm import ( @@ -18,6 +19,7 @@ encode_threat_data, ) + class TestUniqueNames: def test_duplicate_boundary_names_have_different_unique_names(self): random.seed(0) @@ -31,6 +33,7 @@ def test_duplicate_boundary_names_have_different_unique_names(self): assert object_1_uniq_name == "boundary_foo_acf3059e70" assert object_2_uniq_name == "boundary_foo_88f2d9c06f" + class TestAttributes: def test_write_once(self): user = Actor("User") @@ -99,13 +102,17 @@ def test_defaults(self): result_data = Data("Results") result = Dataflow(db, server, "Results", data=result_data, isResponse=True) resp_get_data = Data("HTTP Response") - resp_get = Dataflow(server, user, "HTTP Response", data=resp_get_data, isResponse=True) + resp_get = Dataflow( + server, user, "HTTP Response", data=resp_get_data, isResponse=True + ) test_assumption = Assumption("test assumption") resp_get.assumptions = [test_assumption] req_post_data = Data("JSON") req_post = Dataflow(user, server, "HTTP POST", data=req_post_data) resp_post = Dataflow(server, user, "HTTP Response", isResponse=True) - test_assumption_exclude = Assumption("test assumption", exclude=["ABCD", "BCDE"]) + test_assumption_exclude = Assumption( + "test assumption", exclude=["ABCD", "BCDE"] + ) resp_post.assumptions = [test_assumption_exclude] sql_data = Data("SQL") worker_query = Dataflow(worker, db, "Query", data=sql_data) @@ -115,13 +122,19 @@ def test_defaults(self): assert req_get.srcPort == -1 assert req_get.dstPort == server.port assert req_get.controls.isEncrypted == server.controls.isEncrypted - assert req_get.controls.authenticatesDestination == user.controls.authenticatesDestination + assert ( + req_get.controls.authenticatesDestination + == user.controls.authenticatesDestination + ) assert req_get.protocol == server.protocol assert user.data.issubset(req_get.data) assert server_query.srcPort == -1 assert server_query.dstPort == db.port assert server_query.controls.isEncrypted == db.controls.isEncrypted - assert server_query.controls.authenticatesDestination == server.controls.authenticatesDestination + assert ( + server_query.controls.authenticatesDestination + == server.controls.authenticatesDestination + ) assert server_query.protocol == db.protocol assert server.data.issubset(server_query.data) assert result.srcPort == db.port @@ -141,7 +154,10 @@ def test_defaults(self): assert req_post.srcPort == -1 assert req_post.dstPort == server.port assert req_post.controls.isEncrypted == server.controls.isEncrypted - assert req_post.controls.authenticatesDestination == user.controls.authenticatesDestination + assert ( + req_post.controls.authenticatesDestination + == user.controls.authenticatesDestination + ) assert req_post.protocol == server.protocol assert user.data.issubset(req_post.data) assert resp_post.srcPort == server.port @@ -159,7 +175,10 @@ def test_defaults(self): assert cookie.carriedBy == [req_get, req_post] assert set(cookie.processedBy) == set([user, server]) assert cookie in req_get.data - assert set([d.name for d in req_post.data]) == set([cookie.name, "HTTP", "JSON"]) + assert set([d.name for d in req_post.data]) == set( + [cookie.name, "HTTP", "JSON"] + ) + class TestMethod: def test_defaults(self): @@ -201,7 +220,10 @@ def test_defaults(self): assert tm.check() for case in testCases: t = Threat(SID="", target=default_target, condition=case["condition"]) - assert t.apply(case["target"]), f"Failed to match {case['target']} against {case['condition']}" + assert t.apply(case["target"]), ( + f"Failed to match {case['target']} against {case['condition']}" + ) + class TestFunction: def test_encode_threat_data(self): @@ -222,7 +244,7 @@ def test_encode_threat_data(self): cvss="1.234", response="A test response", assumption=Assumption("Test Assumption", exclude=["INP02"]), - ) + ), ] encoded_findings = encode_threat_data(findings) assert len(encoded_findings) == 2 diff --git a/tests/test_pydantic_models.py b/tests/test_pydantic_models.py index 0942b9e..9c61162 100644 --- a/tests/test_pydantic_models.py +++ b/tests/test_pydantic_models.py @@ -390,6 +390,7 @@ def test_safeset_ignores_invalid_type(self): class TestConditionValidator: def _validate(self, condition: str): import ast + tree = ast.parse(condition, mode="eval") validator = _ConditionValidator(allowed_names=set()) validator.visit(tree) @@ -421,6 +422,7 @@ def test_disallowed_target_method_raises(self): def test_import_node_raises(self): import ast + # Build an Import node manually since parse(..., mode="eval") won't accept it tree = ast.parse("import os", mode="exec") import_node = tree.body[0] diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py index eef1a6f..f223a06 100644 --- a/tests/test_pytmfunc.py +++ b/tests/test_pytmfunc.py @@ -3,14 +3,15 @@ import random import re import tempfile + import pytest from pytm import ( - pytm, + LLM, TM, Action, - Agent, Actor, + Agent, Assumption, Boundary, Classification, @@ -18,27 +19,27 @@ Dataflow, Datastore, ExternalEntity, + Finding, Lambda, - LLM, Lifetime, Process, - Finding, Server, Threat, TLSVersion, loads, + pytm, ) from pytm.pytm import to_serializable with open( os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) + "/pytm/threatlib/threats.json", - "r", ) as threat_file: threats = {t["SID"]: Threat(**t) for t in json.load(threat_file)} output_path = tempfile.gettempdir() + class TestTM: def test_seq(self): random.seed(0) @@ -455,9 +456,17 @@ def test_json_loads(self): "Select", "Response", ] - assert [f.name for f in tm._flows] == ["Request", "Insert", "Select", "Response"] + assert [f.name for f in tm._flows] == [ + "Request", + "Insert", + "Select", + "Response", + ] - assert [d.model_dump(include=["name", "classification", "lifetime"]) for d in tm._data] == [ + assert [ + d.model_dump(include=["name", "classification", "lifetime"]) + for d in tm._data + ] == [ { "name": "Password", "classification": Classification.SECRET, @@ -1868,7 +1877,9 @@ def test_likelihood_copied_from_threat_to_finding(self): TM.reset() tm = TM("test tm", description="aaa") Server("Web Server") - TM._threats = [Threat(SID="T01", target="Server", severity="High", likelihood="Medium")] + TM._threats = [ + Threat(SID="T01", target="Server", severity="High", likelihood="Medium") + ] tm.resolve() server = next(e for e in TM._elements if e.name == "Web Server") @@ -1885,7 +1896,9 @@ def test_override_finding_likelihood_not_overwritten(self): Finding(threat_id="T01", likelihood="High"), ], ) - TM._threats = [Threat(SID="T01", target="Server", severity="High", likelihood="Low")] + TM._threats = [ + Threat(SID="T01", target="Server", severity="High", likelihood="Low") + ] tm.resolve() server = next(e for e in TM._elements if e.name == "Web Server") diff --git a/tests/test_report_util.py b/tests/test_report_util.py index 40dea9f..c4e0259 100644 --- a/tests/test_report_util.py +++ b/tests/test_report_util.py @@ -1,6 +1,6 @@ import pytest -from pytm import Boundary, Finding, Server, TM +from pytm import TM, Boundary, Finding, Server from pytm.report_util import ReportUtils @@ -89,6 +89,7 @@ def test_get_element_type_rejects_non_element(): # Helpers # --------------------------------------------------------------------------- + def _make_finding(element): """Return a Finding attached to *element* using fixed test values.""" return Finding( @@ -109,6 +110,7 @@ def _make_finding(element): # getInScopeFindings — non-element input # --------------------------------------------------------------------------- + def test_get_in_scope_findings_rejects_non_element(): assert ReportUtils.getInScopeFindings(object()) == [] @@ -117,6 +119,7 @@ def test_get_in_scope_findings_rejects_non_element(): # getInScopeFindings — out-of-scope element # --------------------------------------------------------------------------- + def test_get_in_scope_findings_returns_empty_for_out_of_scope_element(): server = Server("OutOfScope") server.inScope = False @@ -131,6 +134,7 @@ def test_get_in_scope_findings_returns_empty_for_out_of_scope_element(): # getInScopeFindings — in-scope element # --------------------------------------------------------------------------- + def test_get_in_scope_findings_returns_findings_for_in_scope_element(): """An in-scope element with findings should have those findings returned.""" server = Server("InScope") @@ -177,6 +181,7 @@ def test_get_in_scope_findings_returns_empty_when_no_findings(): # getInScopeFindings — no cross-element leakage (regression for issue #310) # --------------------------------------------------------------------------- + def test_get_in_scope_findings_does_not_leak_findings_across_elements(): """Findings on one element must not appear on a different element.""" server_a = Server("ServerA") diff --git a/tm.py b/tm.py index 185a0f6..9b403da 100755 --- a/tm.py +++ b/tm.py @@ -1,24 +1,24 @@ #!/usr/bin/env python3 from pytm import ( + LLM, TM, Actor, Agent, + Assumption, Boundary, Classification, Data, Dataflow, Datastore, + DatastoreType, Lambda, - LLM, Server, - DatastoreType, - Assumption, ) tm = TM("my test tm") -tm.description = """This is a sample threat model of a very simple system - a web-based comment system. -The user enters comments and these are added to a database and displayed back to the user. +tm.description = """This is a sample threat model of a very simple system - a web-based comment system. +The user enters comments and these are added to a database and displayed back to the user. The thought is that it is, though simple, a complete enough example to express meaningful threats.""" tm.isOrdered = True tm.mergeResponses = True