add velo artifact results to artifact table (#59)
* add velo artifact results to artifact table * precommit fixes
taylor_socfortress committed
Jul 24, 2023 at 14:35 UTC
a0c9e587ccc1f13e44e74fb800ecbdb2dbc45e54
5 files changed
+81
-1
backend/app/models/artifacts.py
+5
@@ -1,4 +1,7 @@
1
+from datetime import datetime
2
+
3
from sqlalchemy import Column
4
+from sqlalchemy import DateTime
5
from sqlalchemy import Integer
6
from sqlalchemy import String
7
from sqlalchemy.dialects.postgresql import TEXT # Add this line
@@ -18,6 +21,7 @@ class Artifact(db.Model):
21
artifact_name: Column[String] = db.Column(db.String(100))
22
artifact_results: Column[TEXT] = db.Column(TEXT)
23
hostname: Column[String] = db.Column(db.String(100))
24
+ timestamp: Column[DateTime] = db.Column(db.DateTime, default=datetime.utcnow)
25
26
def __init__(self, artifact_name: str, artifact_results: str, hostname: str):
27
"""
@@ -55,6 +59,7 @@ class ArtifactSchema(ma.Schema):
59
"artifact_name",
60
"artifact_results",
61
"hostname",
62
+ "timestamp",
63
)
64
65
backend/app/services/Velociraptor/artifacts.py
+14
@@ -160,6 +160,20 @@ class ArtifactsService:
160
flow_id,
161
artifact,
162
)
163
+
164
+ update_artifact_table_response = self.universal_service.update_artifact_table(
165
+ artifact_name=artifact,
166
+ artifact_results=results,
167
+ hostname=client_id,
168
+ )
169
+
170
+ if not update_artifact_table_response["success"]:
171
+ logger.error(f"Failed to update artifact table: {update_artifact_table_response['message']}")
172
+ return {
173
+ "message": "Failed to update artifact table",
174
+ "success": False,
175
+ }
176
+
177
return results
178
except Exception as err:
179
logger.error(f"Failed to run artifact collection: {err}")
backend/app/services/Velociraptor/universal.py
+30
@@ -6,6 +6,8 @@ import pyvelociraptor
6
from pyvelociraptor import api_pb2
7
from pyvelociraptor import api_pb2_grpc
8
9
+from app import db
10
+from app.models.artifacts import Artifact
11
from app.models.connectors import Connector
12
from app.models.connectors import connector_factory
13
@@ -152,6 +154,34 @@ class UniversalService:
154
vql = f"SELECT * FROM source(client_id='{client_id}', flow_id='{flow_id}', artifact='{artifact}')"
155
return self.execute_query(vql)
156
157
+ def update_artifact_table(self, artifact_name: str, artifact_results: str, hostname: str):
158
+ """
159
+ Update the artifact table with the results of the artifact collection.
160
+
161
+ Args:
162
+ artifact_name (str): The name of the artifact.
163
+ artifact_results (str): The results of the artifact collection.
164
+ hostname (str): The hostname where the artifact was collected.
165
+
166
+ Returns:
167
+ dict: A dictionary with the success status and a message.
168
+ """
169
+ # Covernt the artifact_results from a dict to json
170
+ artifact_results = json.dumps(artifact_results)
171
+ try:
172
+ artifact = Artifact(artifact_name, artifact_results, hostname)
173
+ db.session.add(artifact)
174
+ db.session.commit()
175
+ return {
176
+ "success": True,
177
+ "message": "Successfully updated artifact table",
178
+ }
179
+ except Exception as e:
180
+ return {
181
+ "success": False,
182
+ "message": f"Failed to update artifact table: {e}",
183
+ }
184
+
185
def get_client_id(self, client_name: str):
186
"""
187
Get the client_id associated with a given client_name.
backend/app/static/swagger.json
+1
-1
@@ -1984,7 +1984,7 @@
1984
"description": "Endpoint to get all artifacts for a hostname.",
1985
"parameters": [
1986
{
1987
- "name": "client_name",
1987
+ "name": "hostname",
1988
"in": "path",
1989
"description": "The client name",
1990
"required": true,
backend/migrations/versions/1ec08862d786_add_timestamp_to_artifact_table.py
new
+31
@@ -0,0 +1,31 @@
1
+"""Add timestamp to artifact table
2
+
3
+Revision ID: 1ec08862d786
4
+Revises: 26d24321c4c5
5
+Create Date: 2023-07-24 14:47:39.209785
6
+
7
+"""
8
+import sqlalchemy as sa
9
+from alembic import op
10
+
11
+# revision identifiers, used by Alembic.
12
+revision = "1ec08862d786"
13
+down_revision = "26d24321c4c5"
14
+branch_labels = None
15
+depends_on = None
16
+
17
+
18
+def upgrade():
19
+ # ### commands auto generated by Alembic - please adjust! ###
20
+ with op.batch_alter_table("artifact", schema=None) as batch_op:
21
+ batch_op.add_column(sa.Column("timestamp", sa.DateTime(), nullable=True))
22
+
23
+ # ### end Alembic commands ###
24
+
25
+
26
+def downgrade():
27
+ # ### commands auto generated by Alembic - please adjust! ###
28
+ with op.batch_alter_table("artifact", schema=None) as batch_op:
29
+ batch_op.drop_column("timestamp")
30
+
31
+ # ### end Alembic commands ###