From 0026e00a25f56312a049ce58729fd383a9766fc9 Mon Sep 17 00:00:00 2001 From: Taylor Date: Tue, 11 Jul 2023 07:30:03 -0500 Subject: [PATCH] split models into individual files for modularity --- backend/app/models/artifacts.py | 35 ++++ backend/app/models/cases.py | 34 ++++ backend/app/models/graylog.py | 58 +++++++ backend/app/models/models.py | 255 ++++++++++++++-------------- backend/app/models/wazuh_indexer.py | 50 ++++++ 5 files changed, 304 insertions(+), 128 deletions(-) create mode 100644 backend/app/models/artifacts.py create mode 100644 backend/app/models/cases.py create mode 100644 backend/app/models/graylog.py create mode 100644 backend/app/models/wazuh_indexer.py diff --git a/backend/app/models/artifacts.py b/backend/app/models/artifacts.py new file mode 100644 index 00000000..dec9b7cd --- /dev/null +++ b/backend/app/models/artifacts.py @@ -0,0 +1,35 @@ +from datetime import datetime + +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): + self.artifact_name = artifact_name + self.artifact_results = artifact_results + self.hostname = hostname + + def __repr__(self): + return f"" + + +class ArtifactSchema(ma.Schema): + class Meta: + fields = ( + "id", + "artifact_name", + "artifact_results", + "hostname", + ) + + +artifact_schema = ArtifactSchema() +artifacts_schema = ArtifactSchema(many=True) diff --git a/backend/app/models/cases.py b/backend/app/models/cases.py new file mode 100644 index 00000000..45ba1b3a --- /dev/null +++ b/backend/app/models/cases.py @@ -0,0 +1,34 @@ +from datetime import datetime + +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): + self.case_id = case_id + self.case_name = case_name + self.agents = agents + + def __repr__(self): + return f"" + + +class CaseSchema(ma.Schema): + class Meta: + fields = ( + "id", + "case_id", + "case_name", + "agents", + ) + + +case_schema = CaseSchema() +cases_schema = CaseSchema(many=True) diff --git a/backend/app/models/graylog.py b/backend/app/models/graylog.py new file mode 100644 index 00000000..df10cbea --- /dev/null +++ b/backend/app/models/graylog.py @@ -0,0 +1,58 @@ +from datetime import datetime + +from app import db +from app import ma + +# Class for Graylog allocation which stores throughput metrics +# Generate timestamp for each entry and invoke every 5 minutes. +# Path: backend\app\models.py +class GraylogMetricsAllocation(db.Model): + id = db.Column(db.Integer, primary_key=True) + input_usage = db.Column(db.Float) + output_usage = db.Column(db.Float) + processor_usage = db.Column(db.Float) + input_1_sec_rate = db.Column(db.Float) + output_1_sec_rate = db.Column(db.Float) + total_input = db.Column(db.Float) + total_output = db.Column(db.Float) + timestamp = db.Column(db.DateTime, default=datetime.utcnow) + + def __init__( + self, + input_usage, + output_usage, + processor_usage, + input_1_sec_rate, + output_1_sec_rate, + total_input, + total_output, + ): + self.input_usage = input_usage + self.output_usage = output_usage + self.processor_usage = processor_usage + self.input_1_sec_rate = input_1_sec_rate + self.output_1_sec_rate = output_1_sec_rate + self.total_input = total_input + self.total_output = total_output + + def __repr__(self): + return f"" + + +class GraylogMetricsAllocationSchema(ma.Schema): + class Meta: + fields = ( + "id", + "input_usage", + "output_usage", + "processor_usage", + "input_1_sec_rate", + "output_1_sec_rate", + "total_input", + "total_output", + "timestamp", + ) + + +graylog_metrics_allocation_schema = GraylogMetricsAllocationSchema() +graylog_metrics_allocations_schema = GraylogMetricsAllocationSchema(many=True) diff --git a/backend/app/models/models.py b/backend/app/models/models.py index 3d8c7d73..2d2786fb 100644 --- a/backend/app/models/models.py +++ b/backend/app/models/models.py @@ -1,7 +1,6 @@ from datetime import datetime from loguru import logger -from sqlalchemy.dialects.postgresql import TEXT # Add this line from app import db from app import ma @@ -150,162 +149,162 @@ class Meta: # Class for Wazuh Indexer allocation which stores disk stats and the host. # Generate timestamp for each entry and invoke every 5 minutes. # Path: backend\app\models.py -class WazuhIndexerAllocation(db.Model): - id = db.Column(db.Integer, primary_key=True) - node = db.Column(db.String(100)) - disk_used = db.Column(db.Float) - disk_available = db.Column(db.Float) - disk_total = db.Column(db.Float) - disk_percent = db.Column(db.Float) - timestamp = db.Column(db.DateTime, default=datetime.utcnow) +# class WazuhIndexerAllocation(db.Model): +# id = db.Column(db.Integer, primary_key=True) +# node = db.Column(db.String(100)) +# disk_used = db.Column(db.Float) +# disk_available = db.Column(db.Float) +# disk_total = db.Column(db.Float) +# disk_percent = db.Column(db.Float) +# timestamp = db.Column(db.DateTime, default=datetime.utcnow) - def __init__( - self, - node, - disk_used, - disk_available, - disk_total, - disk_percent, - ): - self.node = node - self.disk_used = disk_used - self.disk_available = disk_available - self.disk_total = disk_total - self.disk_percent = disk_percent +# def __init__( +# self, +# node, +# disk_used, +# disk_available, +# disk_total, +# disk_percent, +# ): +# self.node = node +# self.disk_used = disk_used +# self.disk_available = disk_available +# self.disk_total = disk_total +# self.disk_percent = disk_percent - def __repr__(self): - return f"" +# def __repr__(self): +# return f"" -class WazuhIndexerAllocationSchema(ma.Schema): - class Meta: - fields = ( - "id", - "node", - "disk_used", - "disk_available", - "disk_total", - "disk_percent", - "timestamp", - ) +# class WazuhIndexerAllocationSchema(ma.Schema): +# class Meta: +# fields = ( +# "id", +# "node", +# "disk_used", +# "disk_available", +# "disk_total", +# "disk_percent", +# "timestamp", +# ) -wazuh_indexer_allocation_schema = WazuhIndexerAllocationSchema() -wazuh_indexer_allocations_schema = WazuhIndexerAllocationSchema(many=True) +# wazuh_indexer_allocation_schema = WazuhIndexerAllocationSchema() +# wazuh_indexer_allocations_schema = WazuhIndexerAllocationSchema(many=True) -# Class for Graylog allocation which stores throughput metrics -# Generate timestamp for each entry and invoke every 5 minutes. -# Path: backend\app\models.py -class GraylogMetricsAllocation(db.Model): - id = db.Column(db.Integer, primary_key=True) - input_usage = db.Column(db.Float) - output_usage = db.Column(db.Float) - processor_usage = db.Column(db.Float) - input_1_sec_rate = db.Column(db.Float) - output_1_sec_rate = db.Column(db.Float) - total_input = db.Column(db.Float) - total_output = db.Column(db.Float) - timestamp = db.Column(db.DateTime, default=datetime.utcnow) +# # Class for Graylog allocation which stores throughput metrics +# # Generate timestamp for each entry and invoke every 5 minutes. +# # Path: backend\app\models.py +# class GraylogMetricsAllocation(db.Model): +# id = db.Column(db.Integer, primary_key=True) +# input_usage = db.Column(db.Float) +# output_usage = db.Column(db.Float) +# processor_usage = db.Column(db.Float) +# input_1_sec_rate = db.Column(db.Float) +# output_1_sec_rate = db.Column(db.Float) +# total_input = db.Column(db.Float) +# total_output = db.Column(db.Float) +# timestamp = db.Column(db.DateTime, default=datetime.utcnow) - def __init__( - self, - input_usage, - output_usage, - processor_usage, - input_1_sec_rate, - output_1_sec_rate, - total_input, - total_output, - ): - self.input_usage = input_usage - self.output_usage = output_usage - self.processor_usage = processor_usage - self.input_1_sec_rate = input_1_sec_rate - self.output_1_sec_rate = output_1_sec_rate - self.total_input = total_input - self.total_output = total_output +# def __init__( +# self, +# input_usage, +# output_usage, +# processor_usage, +# input_1_sec_rate, +# output_1_sec_rate, +# total_input, +# total_output, +# ): +# self.input_usage = input_usage +# self.output_usage = output_usage +# self.processor_usage = processor_usage +# self.input_1_sec_rate = input_1_sec_rate +# self.output_1_sec_rate = output_1_sec_rate +# self.total_input = total_input +# self.total_output = total_output - def __repr__(self): - return f"" +# def __repr__(self): +# return f"" -class GraylogMetricsAllocationSchema(ma.Schema): - class Meta: - fields = ( - "id", - "input_usage", - "output_usage", - "processor_usage", - "input_1_sec_rate", - "output_1_sec_rate", - "total_input", - "total_output", - "timestamp", - ) +# class GraylogMetricsAllocationSchema(ma.Schema): +# class Meta: +# fields = ( +# "id", +# "input_usage", +# "output_usage", +# "processor_usage", +# "input_1_sec_rate", +# "output_1_sec_rate", +# "total_input", +# "total_output", +# "timestamp", +# ) -graylog_metrics_allocation_schema = GraylogMetricsAllocationSchema() -graylog_metrics_allocations_schema = GraylogMetricsAllocationSchema(many=True) +# graylog_metrics_allocation_schema = GraylogMetricsAllocationSchema() +# graylog_metrics_allocations_schema = GraylogMetricsAllocationSchema(many=True) # 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)) +# 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): - self.case_id = case_id - self.case_name = case_name - self.agents = agents +# def __init__(self, case_id, case_name, agents): +# self.case_id = case_id +# self.case_name = case_name +# self.agents = agents - def __repr__(self): - return f"" +# def __repr__(self): +# return f"" -class CaseSchema(ma.Schema): - class Meta: - fields = ( - "id", - "case_id", - "case_name", - "agents", - ) +# class CaseSchema(ma.Schema): +# class Meta: +# fields = ( +# "id", +# "case_id", +# "case_name", +# "agents", +# ) -case_schema = CaseSchema() -cases_schema = CaseSchema(many=True) +# case_schema = CaseSchema() +# cases_schema = CaseSchema(many=True) -# 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)) +# # 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): - self.artifact_name = artifact_name - self.artifact_results = artifact_results - self.hostname = hostname +# def __init__(self, artifact_name, artifact_results, hostname): +# self.artifact_name = artifact_name +# self.artifact_results = artifact_results +# self.hostname = hostname - def __repr__(self): - return f"" +# def __repr__(self): +# return f"" -class ArtifactSchema(ma.Schema): - class Meta: - fields = ( - "id", - "artifact_name", - "artifact_results", - "hostname", - ) +# class ArtifactSchema(ma.Schema): +# class Meta: +# fields = ( +# "id", +# "artifact_name", +# "artifact_results", +# "hostname", +# ) -artifact_schema = ArtifactSchema() -artifacts_schema = ArtifactSchema(many=True) +# artifact_schema = ArtifactSchema() +# artifacts_schema = ArtifactSchema(many=True) diff --git a/backend/app/models/wazuh_indexer.py b/backend/app/models/wazuh_indexer.py new file mode 100644 index 00000000..93e184fb --- /dev/null +++ b/backend/app/models/wazuh_indexer.py @@ -0,0 +1,50 @@ +from datetime import datetime + +from app import db +from app import ma + +# Class for Wazuh Indexer allocation which stores disk stats and the host. +# Generate timestamp for each entry and invoke every 5 minutes. +# Path: backend\app\models.py +class WazuhIndexerAllocation(db.Model): + id = db.Column(db.Integer, primary_key=True) + node = db.Column(db.String(100)) + disk_used = db.Column(db.Float) + disk_available = db.Column(db.Float) + disk_total = db.Column(db.Float) + disk_percent = db.Column(db.Float) + timestamp = db.Column(db.DateTime, default=datetime.utcnow) + + def __init__( + self, + node, + disk_used, + disk_available, + disk_total, + disk_percent, + ): + self.node = node + self.disk_used = disk_used + self.disk_available = disk_available + self.disk_total = disk_total + self.disk_percent = disk_percent + + def __repr__(self): + return f"" + + +class WazuhIndexerAllocationSchema(ma.Schema): + class Meta: + fields = ( + "id", + "node", + "disk_used", + "disk_available", + "disk_total", + "disk_percent", + "timestamp", + ) + + +wazuh_indexer_allocation_schema = WazuhIndexerAllocationSchema() +wazuh_indexer_allocations_schema = WazuhIndexerAllocationSchema(many=True)