Conversation
colesmj
left a comment
There was a problem hiding this comment.
Thank you for submitting this refactoring of the pytm base code; it was long overdue. The change to pydantic instead of the legacy var mechanism will make for cleaner code and also create some new possibilities for extending pytm.
No threats needed to change to support this refactor (e.g. the condition evals were all "safe" with the checks you added)?
| class Lambda(Asset): | ||
| """A lambda function running in a Function-as-a-Service (FaaS) environment.""" | ||
|
|
||
| onAWS: bool = Field(default=True, description="Is this lambda on AWS") |
There was a problem hiding this comment.
Just want to verify - did you intend to recreate this Field in the subclass vs setting the Field inherited?
There was a problem hiding this comment.
Was just following the original code where onAWS defaults to True for Lambda. To be honest I'm not a fan of using "onAWS" or calling this "Lambda" but I didn't want to change the API for now.
| default=None, | ||
| description="Is a response to this data flow" | ||
| ) | ||
| srcPort: int = Field(default=-1, description="Source TCP port") |
There was a problem hiding this comment.
A more general comment, not specific to this review but convenient: The srcPort and dstPort descriptions should drop 'TCP' to give them some flexibility to be dataflows for more than just TCP networks.
There was a problem hiding this comment.
I agree, especially when we have a protocol field here too. What are your thoughts about setting the protocol field to default to TCP?
There was a problem hiding this comment.
From a threat perspective that is probably less useful than setting it to HTTP or HTTPS. In truth, any data flow may have layered protocols (e.g. TLS over HTTP over TCP over IPv4); obviously some are more informative for threat analysis than others.
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| class var(object): |
Thanks for the review and the comments! There are a bunch of changes I'd like to make but I wanted to keep it invisible to most end users. Maybe with more drastic changes this might become part of a 2.0 release? No threats needed to change for now. I did think about moving the threat library to python code, but there was not a lot to gain other than using enums and classes (to save string matching). |
Replaces the custom var descriptor system with Pydantic v2 BaseModel. This simplifies validation and makes it easier to extend. The public API is unchanged, existing threat model scripts should work without modification.
Changes:
Bug fixes:
Added Tests:
Added tests/test_pydantic_models.py covering the new validation and coercion behaviour (DataSet, field validators, Threat legacy field mapping). All existing tests pass.