From d7f1f93093e3a9c119726dbf5c3c312166c6999d Mon Sep 17 00:00:00 2001
From: erev0s <erev0s@mail.com>
Date: Mon, 25 Nov 2024 14:58:24 +0200
Subject: [PATCH] custom generic error

---
 config.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/config.py b/config.py
index a3f4f7ff..17e65c5e 100644
--- a/config.py
+++ b/config.py
@@ -2,6 +2,7 @@
 import connexion
 from flask import jsonify
 from flask_sqlalchemy import SQLAlchemy
+from connexion.exceptions import ProblemException
 
 vuln_app = connexion.App(__name__, specification_dir='./openapi_specs')
 
@@ -13,13 +14,14 @@
 # start the db
 db = SQLAlchemy(vuln_app.app)
 
-
-@vuln_app.app.errorhandler(401)
-def custom_401(error):
-    # Custom 401 to match the original response sent by Vampi
-    response = jsonify({"status": "fail", "message": "Invalid token. Please log in again."})
-    response.status_code = 401
+def custom_problem_handler(error):
+    # Custom error handler for clarity in structure
+    response = jsonify({
+        "status": "fail",
+        "message": getattr(error, "detail", "An error occurred"),
+    })
+    response.status_code = error.status
     return response
-
+vuln_app.add_error_handler(ProblemException, custom_problem_handler)
 
 vuln_app.add_api('openapi3.yml')