-
-
Notifications
You must be signed in to change notification settings - Fork 41
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
added docstrings and typehints to models (#3)
- Loading branch information
1 parent
9db52fe
commit f6d3a87
Showing
9 changed files
with
444 additions
and
160 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,35 +1,62 @@ | ||
from datetime import datetime | ||
from sqlalchemy import Column | ||
from sqlalchemy import Integer | ||
from sqlalchemy import String | ||
from sqlalchemy.dialects.postgresql import TEXT # Add this line | ||
|
||
from app import db | ||
from app import ma | ||
from sqlalchemy.dialects.postgresql import TEXT # Add this line | ||
|
||
# Class for artifacts collected which stores the artifact name, artificat results (json), hostname | ||
|
||
# Path: backend\app\models.py | ||
class Artifact(db.Model): | ||
id = db.Column(db.Integer, primary_key=True) | ||
artifact_name = db.Column(db.String(100)) | ||
artifact_results = db.Column(TEXT) | ||
hostname = db.Column(db.String(100)) | ||
|
||
def __init__(self, artifact_name, artifact_results, hostname): | ||
""" | ||
Class for artifacts collected which stores the artifact name, artifact results (json), and hostname. | ||
This class inherits from SQLAlchemy's Model class. | ||
""" | ||
|
||
id: Column[Integer] = db.Column(db.Integer, primary_key=True) | ||
artifact_name: Column[String] = db.Column(db.String(100)) | ||
artifact_results: Column[TEXT] = db.Column(TEXT) | ||
hostname: Column[String] = db.Column(db.String(100)) | ||
|
||
def __init__(self, artifact_name: str, artifact_results: str, hostname: str): | ||
""" | ||
Initialize a new instance of the Artifact class. | ||
:param artifact_name: The name of the artifact. | ||
:param artifact_results: The results of the artifact, stored as a JSON string. | ||
:param hostname: The hostname where the artifact was collected. | ||
""" | ||
self.artifact_name = artifact_name | ||
self.artifact_results = artifact_results | ||
self.hostname = hostname | ||
|
||
def __repr__(self): | ||
def __repr__(self) -> str: | ||
""" | ||
Returns a string representation of the Artifact instance. | ||
:return: A string representation of the artifact name. | ||
""" | ||
return f"<Artifact {self.artifact_name}>" | ||
|
||
|
||
class ArtifactSchema(ma.Schema): | ||
""" | ||
Schema for serializing and deserializing instances of the Artifact class. | ||
""" | ||
|
||
class Meta: | ||
fields = ( | ||
""" | ||
Meta class defines the fields to be serialized/deserialized. | ||
""" | ||
|
||
fields: tuple = ( | ||
"id", | ||
"artifact_name", | ||
"artifact_results", | ||
"hostname", | ||
) | ||
|
||
|
||
artifact_schema = ArtifactSchema() | ||
artifacts_schema = ArtifactSchema(many=True) | ||
artifact_schema: ArtifactSchema = ArtifactSchema() | ||
artifacts_schema: ArtifactSchema = ArtifactSchema(many=True) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,34 +1,61 @@ | ||
from datetime import datetime | ||
from sqlalchemy import Column | ||
from sqlalchemy import Integer | ||
from sqlalchemy import String | ||
|
||
from app import db | ||
from app import ma | ||
|
||
# Class for cases which stores the case ID, case name, list of agents | ||
|
||
# Path: backend\app\models.py | ||
class Case(db.Model): | ||
id = db.Column(db.Integer, primary_key=True) | ||
case_id = db.Column(db.Integer) | ||
case_name = db.Column(db.String(100)) | ||
agents = db.Column(db.String(1000)) | ||
|
||
def __init__(self, case_id, case_name, agents): | ||
""" | ||
Class for cases which stores the case ID, case name, and a list of agents. | ||
This class inherits from SQLAlchemy's Model class. | ||
""" | ||
|
||
id: Column[Integer] = db.Column(db.Integer, primary_key=True) | ||
case_id: Column[Integer] = db.Column(db.Integer) | ||
case_name: Column[String] = db.Column(db.String(100)) | ||
agents: Column[String] = db.Column(db.String(1000)) | ||
|
||
def __init__(self, case_id: int, case_name: str, agents: str): | ||
""" | ||
Initialize a new instance of the Case class. | ||
:param case_id: The ID of the case. | ||
:param case_name: The name of the case. | ||
:param agents: A comma-separated string of agents associated with the case. | ||
""" | ||
self.case_id = case_id | ||
self.case_name = case_name | ||
self.agents = agents | ||
|
||
def __repr__(self): | ||
def __repr__(self) -> str: | ||
""" | ||
Returns a string representation of the Case instance. | ||
:return: A string representation of the case ID. | ||
""" | ||
return f"<Case {self.case_id}>" | ||
|
||
|
||
class CaseSchema(ma.Schema): | ||
""" | ||
Schema for serializing and deserializing instances of the Case class. | ||
""" | ||
|
||
class Meta: | ||
fields = ( | ||
""" | ||
Meta class defines the fields to be serialized/deserialized. | ||
""" | ||
|
||
fields: tuple = ( | ||
"id", | ||
"case_id", | ||
"case_name", | ||
"agents", | ||
) | ||
|
||
|
||
case_schema = CaseSchema() | ||
cases_schema = CaseSchema(many=True) | ||
case_schema: CaseSchema = CaseSchema() | ||
cases_schema: CaseSchema = CaseSchema(many=True) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.