Auth (#98)
* auth fields to user model * required libraries * login and register route * jwt creation and validation * precommit fixes
taylor_socfortress committed
Sep 15, 2023 at 12:24 UTC
dc225c1f28dc99d9d12f75406d7c67d884866fbf
9 files changed
+390
backend/app/__init__.py
+4
@@ -1,11 +1,13 @@
1
from flask import Flask
2
from flask_cors import CORS
3
+from flask_jwt_extended import JWTManager
4
from flask_marshmallow import Marshmallow
5
from flask_migrate import Migrate
6
from flask_sqlalchemy import SQLAlchemy
7
from flask_swagger_ui import get_swaggerui_blueprint
8
9
app = Flask(__name__)
10
+jwt = JWTManager(app)
11
12
SWAGGER_URL = "/api/docs" # URL for exposing Swagger UI (without trailing '/')
13
API_URL = "/static/swagger.json" # Our API url (can of course be a local resource)
@@ -63,6 +65,7 @@ from app.routes.shuffle import bp as shuffle_bp
65
from app.routes.smtp import bp as smtp_bp
66
from app.routes.sublime import bp as sublime_bp
67
from app.routes.threatintel import bp as threatintel_bp
68
+from app.routes.users import bp as users_bp
69
from app.routes.velociraptor import bp as velociraptor_bp
70
from app.routes.wazuhindexer import bp as wazuhindexer_bp
71
@@ -83,3 +86,4 @@ app.register_blueprint(threatintel_bp) # Register the threatintel blueprint
86
app.register_blueprint(customers_bp) # Register the customers blueprint
87
app.register_blueprint(dnstwist_bp) # Register the dnstwist blueprint
88
app.register_blueprint(cortex_bp) # Register the cortex blueprint
89
+app.register_blueprint(users_bp) # Register the login blueprint
backend/app/models/users.py
+8
@@ -24,10 +24,14 @@ class Users(db.Model):
24
imageFile: Column[String] = db.Column(db.String(64))
25
notifications: Column[Integer] = db.Column(db.SmallInteger, nullable=False)
26
createdAt: Column[DateTime] = db.Column(db.DateTime, default=datetime.utcnow)
27
+ passwordHash = db.Column(db.String(128), nullable=False)
28
+ jwtToken = db.Column(db.String(512))
29
+ tokenExpiry = db.Column(db.DateTime)
30
31
def __init__(
32
self,
33
customerCode: str,
34
+ passwordHash: str,
35
notifications: int,
36
usersFirstName: str = None,
37
usersLastName: str = None,
@@ -56,6 +60,7 @@ class Users(db.Model):
60
self.imageFile = imageFile
61
self.notifications = notifications
62
self.createdAt = createdAt
63
+ self.passwordHash = passwordHash
64
65
def __repr__(self) -> str:
66
"""
@@ -86,6 +91,9 @@ class UsersSchema(ma.Schema):
91
"imageFile",
92
"notifications",
93
"createdAt",
94
+ "passwordHash",
95
+ "jwtToken",
96
+ "tokenExpiry",
97
)
98
99
backend/app/routes/agents.py
+12
@@ -2,11 +2,14 @@ from typing import Any
2
3
from flask import Blueprint
4
from flask import jsonify
5
+from flask import request
6
+from flask_jwt_extended import verify_jwt_in_request
7
from loguru import logger
8
9
from app.models.agents import agent_metadata_schema
10
from app.services.agents.agents import AgentService
11
from app.services.agents.agents import AgentSyncService
12
+from app.services.validator import check_jwt_in_db
13
from app.services.wazuh_manager.agent import WazuhManagerAgentService
14
from app.services.wazuh_manager.universal import UniversalService
15
from app.services.wazuh_manager.vulnerability import VulnerabilityService
@@ -14,6 +17,15 @@ from app.services.wazuh_manager.vulnerability import VulnerabilityService
17
bp = Blueprint("agents", __name__)
18
19
20
+@bp.before_request
21
+def before_request_func():
22
+ if request.endpoint != "agents.get_agents": # Example of endpoint that doesn't require JWT
23
+ verify_jwt_in_request()
24
+ response = check_jwt_in_db()
25
+ if response:
26
+ return response
27
+
28
+
29
@bp.route("/agents", methods=["GET"])
30
def get_agents() -> Any:
31
"""
backend/app/routes/users.py
new
+25
@@ -0,0 +1,25 @@
1
+# routes/users.py
2
+
3
+from flask import Blueprint
4
+from flask import jsonify
5
+from flask import request
6
+
7
+from app.services.users.universal import UniversalService
8
+
9
+bp = Blueprint("users", __name__)
10
+
11
+
12
+@bp.route("/register", methods=["POST"])
13
+def register():
14
+ if not request.is_json:
15
+ return jsonify({"message": "Missing JSON in request", "success": False}), 400
16
+ response, status_code = UniversalService.register_user(request.json)
17
+ return jsonify(response), status_code
18
+
19
+
20
+@bp.route("/login", methods=["POST"])
21
+def login():
22
+ if not request.is_json:
23
+ return jsonify({"message": "Missing JSON in request", "success": False}), 400
24
+ response, status_code = UniversalService.login_user(request.json)
25
+ return jsonify(response), status_code
backend/app/services/users/universal.py
new
+92
@@ -0,0 +1,92 @@
1
+from datetime import datetime
2
+from datetime import timedelta
3
+
4
+import bcrypt
5
+from flask_jwt_extended import create_access_token
6
+
7
+from app import db
8
+from app.models.users import Users
9
+
10
+
11
+class UniversalService:
12
+ @staticmethod
13
+ def validate_user_input(data):
14
+ required_fields = ["customerCode", "usersEmail", "password"]
15
+ for field in required_fields:
16
+ if not data.get(field):
17
+ return False
18
+ return True
19
+
20
+ @staticmethod
21
+ def user_exists(email):
22
+ return Users.query.filter_by(usersEmail=email).first()
23
+
24
+ @staticmethod
25
+ def hash_password(password):
26
+ return bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")
27
+
28
+ @staticmethod
29
+ def create_user(data, hashed_password):
30
+ new_user = Users(
31
+ customerCode=data["customerCode"],
32
+ usersFirstName=data.get("usersFirstName"),
33
+ usersLastName=data.get("usersLastName"),
34
+ usersEmail=data["usersEmail"],
35
+ usersRole=data.get("usersRole"),
36
+ imageFile=data.get("imageFile"),
37
+ notifications=data.get("notifications", 0),
38
+ passwordHash=hashed_password,
39
+ )
40
+ db.session.add(new_user)
41
+ db.session.commit()
42
+
43
+ @staticmethod
44
+ def register_user(data):
45
+ if not UniversalService.validate_user_input(data):
46
+ return {"message": "Missing mandatory fields", "success": False}, 400
47
+
48
+ if UniversalService.user_exists(data["usersEmail"]):
49
+ return {"message": "User already exists", "success": False}, 409
50
+
51
+ hashed_password = UniversalService.hash_password(data["password"])
52
+ UniversalService.create_user(data, hashed_password)
53
+
54
+ return {"message": "User registered successfully", "success": True}, 201
55
+
56
+ @staticmethod
57
+ def authenticate_user(email, password):
58
+ user = UniversalService.user_exists(email)
59
+ if not user:
60
+ return {"message": "User not found", "success": False}, 404
61
+
62
+ if not bcrypt.checkpw(password.encode("utf-8"), user.passwordHash.encode("utf-8")):
63
+ return {"message": "Incorrect password", "success": False}, 401
64
+
65
+ return user
66
+
67
+ @staticmethod
68
+ def generate_token(email):
69
+ return create_access_token(identity=email, expires_delta=timedelta(days=1))
70
+
71
+ @staticmethod
72
+ def login_user(data):
73
+ email = data.get("email")
74
+ password = data.get("password")
75
+
76
+ if not email or not password:
77
+ return {"message": "Missing email or password parameter", "success": False}, 400
78
+
79
+ user = UniversalService.authenticate_user(email, password)
80
+ if isinstance(user, tuple):
81
+ return user
82
+
83
+ access_token = UniversalService.generate_token(email)
84
+
85
+ user.jwtToken = access_token
86
+ if user.tokenExpiry is None:
87
+ user.tokenExpiry = datetime.utcnow()
88
+ user.tokenExpiry += timedelta(days=1)
89
+
90
+ db.session.commit()
91
+
92
+ return {"message": "User logged in successfully", "success": True, "token": access_token}, 200
backend/app/services/validator.py
new
+44
@@ -0,0 +1,44 @@
1
+from functools import wraps
2
+
3
+from flask import jsonify
4
+from flask_jwt_extended import decode_token
5
+from flask_jwt_extended import get_jwt
6
+from flask_jwt_extended import get_jwt_identity
7
+from flask_jwt_extended import jwt_required
8
+
9
+from app.models.users import Users
10
+
11
+
12
+def check_jwt_in_db():
13
+ """Check if JWT exists in the database and is valid."""
14
+ current_user_identity = get_jwt_identity() # Get identity email from current token
15
+ user = Users.query.filter_by(usersEmail=current_user_identity).first() # Fetch user from DB
16
+
17
+ if not user:
18
+ return jsonify({"message": "User not found", "success": False}), 404
19
+
20
+ # Extract the token from the current request
21
+ current_token = get_jwt()["jti"] # 'jti' is the unique identifier for a JWT
22
+
23
+ # Decode the stored token to get its JTI
24
+ decoded_stored_token = decode_token(user.jwtToken)
25
+ stored_jti = decoded_stored_token["jti"]
26
+
27
+ if stored_jti != current_token:
28
+ return jsonify({"message": "Token mismatch", "success": False}), 401
29
+
30
+ return None # Return None if everything is fine
31
+
32
+
33
+def jwt_db_check(fn):
34
+ """Decorator to check JWT against the database."""
35
+
36
+ @wraps(fn)
37
+ @jwt_required() # First ensure that a valid JWT token is present
38
+ def wrapper(*args, **kwargs):
39
+ response = check_jwt_in_db()
40
+ if response:
41
+ return response
42
+ return fn(*args, **kwargs)
43
+
44
+ return wrapper
backend/app/static/swagger.json
+168
@@ -140,9 +140,177 @@
140
"description": "Find out more",
141
"url": "http://swagger.io"
142
}
143
+ },
144
+ {
145
+ "name": "Authenticator",
146
+ "description": "Everything about Authenticator",
147
+ "externalDocs": {
148
+ "description": "Find out more",
149
+ "url": "http://swagger.io"
150
+ }
151
}
152
],
153
"paths": {
154
+ "/register": {
155
+ "post": {
156
+ "tags": ["Authenticator"],
157
+ "summary": "Register a new user",
158
+ "description": "Endpoint to register a new user.",
159
+ "requestBody": {
160
+ "content": {
161
+ "application/json": {
162
+ "schema": {
163
+ "type": "object",
164
+ "properties": {
165
+ "customerCode": {
166
+ "type": "string",
167
+ "description": "The code of the customer"
168
+ },
169
+ "usersFirstName": {
170
+ "type": "string",
171
+ "description": "The first name of the user"
172
+ },
173
+ "usersLastName": {
174
+ "type": "string",
175
+ "description": "The last name of the user"
176
+ },
177
+ "usersEmail": {
178
+ "type": "string",
179
+ "description": "The email of the user"
180
+ },
181
+ "usersRole": {
182
+ "type": "string",
183
+ "description": "The role of the user"
184
+ },
185
+ "imageFile": {
186
+ "type": "string",
187
+ "description": "The file name of the user's image"
188
+ },
189
+ "notifications": {
190
+ "type": "integer",
191
+ "description": "Whether the user has notifications enabled (1) or not (0)"
192
+ },
193
+ "password": {
194
+ "type": "string",
195
+ "description": "The password of the user"
196
+ }
197
+ },
198
+ "required": ["customerCode", "usersEmail", "password"]
199
+ }
200
+ }
201
+ }
202
+ },
203
+ "responses": {
204
+ "201": {
205
+ "description": "User registered successfully.",
206
+ "content": {
207
+ "application/json": {
208
+ "schema": {
209
+ "type": "object",
210
+ "properties": {
211
+ "message": {
212
+ "type": "string"
213
+ }
214
+ }
215
+ }
216
+ }
217
+ }
218
+ },
219
+ "400": {
220
+ "description": "Bad request, missing or invalid parameters",
221
+ "content": {
222
+ "application/json": {
223
+ "schema": {
224
+ "type": "object",
225
+ "properties": {
226
+ "error": {
227
+ "type": "string"
228
+ }
229
+ }
230
+ }
231
+ }
232
+ }
233
+ },
234
+ "409": {
235
+ "description": "Conflict, user already exists",
236
+ "content": {
237
+ "application/json": {
238
+ "schema": {
239
+ "type": "object",
240
+ "properties": {
241
+ "error": {
242
+ "type": "string"
243
+ }
244
+ }
245
+ }
246
+ }
247
+ }
248
+ }
249
+ }
250
+ }
251
+ },
252
+ "/login": {
253
+ "post": {
254
+ "tags": ["Authenticator"],
255
+ "summary": "Run login attempt",
256
+ "description": "Endpoint to login.",
257
+ "requestBody": {
258
+ "content": {
259
+ "application/json": {
260
+ "schema": {
261
+ "type": "object",
262
+ "properties": {
263
+ "email": {
264
+ "type": "string",
265
+ "description": "User's email"
266
+ },
267
+ "password": {
268
+ "type": "string",
269
+ "description": "User's password"
270
+ }
271
+ },
272
+ "required": ["email", "password"]
273
+ }
274
+ }
275
+ }
276
+ },
277
+ "responses": {
278
+ "200": {
279
+ "description": "Login successful.",
280
+ "content": {
281
+ "application/json": {
282
+ "schema": {
283
+ "type": "object",
284
+ "properties": {
285
+ "data": {
286
+ "type": "array",
287
+ "items": {
288
+ "type": "string"
289
+ }
290
+ }
291
+ }
292
+ }
293
+ }
294
+ }
295
+ },
296
+ "400": {
297
+ "description": "Invalid input",
298
+ "content": {
299
+ "application/json": {
300
+ "schema": {
301
+ "type": "object",
302
+ "properties": {
303
+ "error": {
304
+ "type": "string"
305
+ }
306
+ }
307
+ }
308
+ }
309
+ }
310
+ }
311
+ }
312
+ }
313
+ },
314
"/cortex/analyzers": {
315
"get": {
316
"tags": ["Cortex"],
backend/migrations/versions/92ef2c136375_add_auth_fields.py
new
+35
@@ -0,0 +1,35 @@
1
+"""Add Auth fields
2
+
3
+Revision ID: 92ef2c136375
4
+Revises: c8c903153246
5
+Create Date: 2023-09-15 10:26:42.478296
6
+
7
+"""
8
+import sqlalchemy as sa
9
+from alembic import op
10
+
11
+# revision identifiers, used by Alembic.
12
+revision = "92ef2c136375"
13
+down_revision = "c8c903153246"
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("users", schema=None) as batch_op:
21
+ batch_op.add_column(sa.Column("passwordHash", sa.String(length=128), nullable=False))
22
+ batch_op.add_column(sa.Column("jwtToken", sa.String(length=512), nullable=True))
23
+ batch_op.add_column(sa.Column("tokenExpiry", sa.DateTime(), nullable=True))
24
+
25
+ # ### end Alembic commands ###
26
+
27
+
28
+def downgrade():
29
+ # ### commands auto generated by Alembic - please adjust! ###
30
+ with op.batch_alter_table("users", schema=None) as batch_op:
31
+ batch_op.drop_column("tokenExpiry")
32
+ batch_op.drop_column("jwtToken")
33
+ batch_op.drop_column("passwordHash")
34
+
35
+ # ### end Alembic commands ###
backend/requirements.in
+2
@@ -1,3 +1,4 @@
1
+bcrypt
2
blueprint
3
cortex4py
4
dfir_iris_client
@@ -6,6 +7,7 @@ elasticsearch7==7.10.1
7
environs
8
flask
9
flask-cors
10
+Flask-JWT-Extended
11
flask-marshmallow
12
flask-migrate
13
flask-sqlalchemy