main
py 248 lines 9.42 KB
Raw
1 import os
2 from contextlib import asynccontextmanager
3 from pathlib import Path
4
5 from dotenv import load_dotenv
6
7 # Load .env into os.environ before any `from app...` import. AuthHandler's class body
8 # (app/auth/utils.py) calls _load_jwt_secret() at import time and reads JWT_SECRET
9 # directly from os.environ — environs >= 14 no longer populates os.environ from
10 # read_env(), so without this line the backend refuses to boot unless JWT_SECRET is
11 # already exported in the shell (or uvicorn was invoked with --env-file).
12 load_dotenv(Path(__file__).parent.parent / ".env")
13
14 import uvicorn # noqa: E402
15 from fastapi import APIRouter # noqa: E402
16 from fastapi import FastAPI # noqa: E402
17 from fastapi import HTTPException # noqa: E402
18 from fastapi.exceptions import RequestValidationError # noqa: E402
19 from fastapi.middleware.cors import CORSMiddleware # noqa: E402
20 from fastapi.staticfiles import StaticFiles # noqa: E402
21 from loguru import logger # noqa: E402
22
23 from app.auth.utils import AuthHandler # noqa: E402
24 from app.data_store.data_store_setup import create_buckets
25 from app.db.db_session import SQLALCHEMY_DATABASE_URI_NO_DB
26 from app.db.db_session import async_engine
27 from app.db.db_setup import add_connectors
28 from app.db.db_setup import apply_migrations
29 from app.db.db_setup import create_available_integrations
30 from app.db.db_setup import create_available_network_connectors
31 from app.db.db_setup import create_copilot_user_if_not_exists
32 from app.db.db_setup import create_database_if_not_exists
33 from app.db.db_setup import create_roles
34 from app.db.db_setup import delete_connectors
35 from app.db.db_setup import ensure_admin_user
36 from app.db.db_setup import ensure_scheduler_user
37 from app.db.db_setup import ensure_scheduler_user_removed
38 from app.middleware.exception_handlers import custom_http_exception_handler
39 from app.middleware.exception_handlers import validation_exception_handler
40 from app.middleware.exception_handlers import value_error_handler
41
42 # from app.routers import ask_socfortress
43 from app.routers import active_response
44 from app.routers import agents
45 from app.routers import ai_analyst
46 from app.routers import alert_creation_settings
47 from app.routers import auth
48 from app.routers import bitdefender
49 from app.routers import carbonblack
50 from app.routers import cato
51 from app.routers import connectors
52 from app.routers import copilot_action
53 from app.routers import copilot_mcp
54 from app.routers import copilot_searches
55 from app.routers import cortex
56 from app.routers import crowdstrike
57 from app.routers import customer_portal
58 from app.routers import customer_provisioning
59 from app.routers import customers
60 from app.routers import darktrace
61 from app.routers import data_store
62 from app.routers import defenderforendpoint
63 from app.routers import duo
64 from app.routers import github_audit
65 from app.routers import grafana
66 from app.routers import graylog
67 from app.routers import healthcheck
68 from app.routers import huntress
69 from app.routers import incidents
70 from app.routers import influxdb
71 from app.routers import integrations
72 from app.routers import license
73 from app.routers import logs
74 from app.routers import microsoft_patch_tuesday
75 from app.routers import mimecast
76 from app.routers import modules
77 from app.routers import monitoring_alert
78 from app.routers import network_connectors
79 from app.routers import notifications
80 from app.routers import nuclei
81 from app.routers import office365
82 from app.routers import portainer
83 from app.routers import sap_siem
84 from app.routers import scheduler
85 from app.routers import scoutsuite
86 from app.routers import shuffle
87 from app.routers import siem
88 from app.routers import socfortress_mdr
89 from app.routers import stack_provisioning
90 from app.routers import sublime
91 from app.routers import talon
92 from app.routers import threat_intel
93 from app.routers import velociraptor
94 from app.routers import version
95 from app.routers import wazuh_indexer
96 from app.routers import wazuh_manager
97 from app.schedulers.scheduler import get_scheduler_instance
98 from app.schedulers.scheduler import init_scheduler
99
100 # from app.middleware.logger import log_requests
101
102
103 auth_handler = AuthHandler()
104 server_ip = os.getenv("SERVER_IP", "localhost")
105 environment = os.getenv("ENVIRONMENT", "PRODUCTION")
106
107 # ! Not needed for now maybe revist later ! #
108 # ssl_keyfile = os.path.join(os.path.dirname(__file__), "../nginx/server.key")
109 # ssl_certfile = os.path.join(os.path.dirname(__file__), "../nginx/server.crt")
110
111
112 @asynccontextmanager
113 async def lifespan(_app: FastAPI):
114 # ── startup ──
115 logger.info("Initializing database")
116 if environment == "PRODUCTION":
117 await create_database_if_not_exists(db_url=SQLALCHEMY_DATABASE_URI_NO_DB, db_name="copilot")
118 await create_copilot_user_if_not_exists(db_url=SQLALCHEMY_DATABASE_URI_NO_DB, db_user_name="copilot")
119 apply_migrations()
120 await create_buckets()
121 await add_connectors(async_engine)
122 await delete_connectors(async_engine)
123 await create_roles(async_engine)
124 await create_available_integrations(async_engine)
125 await create_available_network_connectors(async_engine)
126 await ensure_admin_user(async_engine)
127 await ensure_scheduler_user(async_engine)
128
129 # Initialize the scheduler
130 scheduler = await init_scheduler()
131 if not scheduler.running:
132 logger.info("Scheduler is not running, starting now...")
133 scheduler.start()
134
135 yield
136
137 # ── shutdown ──
138 logger.info("Shutting down scheduler")
139 scheduler = await get_scheduler_instance()
140 if scheduler.running:
141 logger.info("Scheduler is running, shutting down now...")
142 scheduler.shutdown()
143 await ensure_scheduler_user_removed(async_engine)
144
145
146 app = FastAPI(description="CoPilot API", version="0.1.0", title="CoPilot API", lifespan=lifespan)
147
148 # Create an APIRouter with a prefix of `/api`
149 api_router = APIRouter(prefix="/api")
150
151
152 # Allow all origins, methods and headers
153 app.add_middleware(
154 CORSMiddleware,
155 allow_origins=["*"],
156 allow_credentials=True,
157 allow_methods=["*"],
158 allow_headers=["*"],
159 )
160
161
162 ################## ! Middleware LOGGING TO `log_entry` table ! ##################
163 # Comment out logging for now, not sure I want to use it
164 # app.middleware("http")(log_requests) # using the imported middleware
165
166
167 ################## ! Exception Handlers ! ##################
168 app.add_exception_handler(HTTPException, custom_http_exception_handler)
169 app.add_exception_handler(RequestValidationError, validation_exception_handler)
170 app.add_exception_handler(ValueError, value_error_handler)
171
172
173 ################## ! INCLUDE ROUTES ! ##################
174 api_router.include_router(connectors.router)
175 api_router.include_router(wazuh_indexer.router)
176 api_router.include_router(auth.router)
177 api_router.include_router(cato.router)
178 api_router.include_router(wazuh_manager.router)
179 api_router.include_router(agents.router)
180 api_router.include_router(graylog.router)
181 api_router.include_router(cortex.router)
182 api_router.include_router(velociraptor.router)
183 api_router.include_router(shuffle.router)
184 api_router.include_router(github_audit.router)
185 api_router.include_router(sublime.router)
186 api_router.include_router(microsoft_patch_tuesday.router)
187 api_router.include_router(customers.router)
188 api_router.include_router(healthcheck.router)
189 api_router.include_router(logs.router)
190 api_router.include_router(influxdb.router)
191 api_router.include_router(version.router)
192 api_router.include_router(grafana.router)
193 api_router.include_router(customer_provisioning.router)
194 api_router.include_router(threat_intel.router)
195 # ! Commenting out for now, will revist later if needed ! #
196 # api_router.include_router(ask_socfortress.router)
197 api_router.include_router(alert_creation_settings.router)
198 api_router.include_router(integrations.router)
199 api_router.include_router(office365.router)
200 api_router.include_router(copilot_action.router)
201 api_router.include_router(copilot_mcp.router)
202 api_router.include_router(mimecast.router)
203 api_router.include_router(copilot_searches.router)
204 api_router.include_router(scheduler.router)
205 api_router.include_router(monitoring_alert.router)
206 api_router.include_router(sap_siem.router)
207 api_router.include_router(stack_provisioning.router)
208 api_router.include_router(data_store.router)
209 api_router.include_router(customer_portal.router)
210 api_router.include_router(active_response.router)
211 api_router.include_router(huntress.router)
212 api_router.include_router(license.router)
213 api_router.include_router(modules.router)
214 api_router.include_router(carbonblack.router)
215 api_router.include_router(network_connectors.router)
216 api_router.include_router(crowdstrike.router)
217 api_router.include_router(bitdefender.router)
218 api_router.include_router(socfortress_mdr.router)
219 api_router.include_router(scoutsuite.router)
220 api_router.include_router(nuclei.router)
221 api_router.include_router(duo.router)
222 api_router.include_router(portainer.router)
223 api_router.include_router(incidents.router)
224 api_router.include_router(ai_analyst.router)
225 api_router.include_router(notifications.router)
226 api_router.include_router(darktrace.router)
227 api_router.include_router(defenderforendpoint.router)
228 api_router.include_router(siem.router)
229 api_router.include_router(talon.router)
230
231 # Include the APIRouter in the FastAPI app
232 app.include_router(api_router)
233
234
235 # Create `scoutsuite-report` directory if it doesnt exist
236 if not os.path.exists("scoutsuite-report"):
237 os.makedirs("scoutsuite-report")
238
239 app.mount("/scoutsuite-report", StaticFiles(directory="scoutsuite-report"), name="scoutsuite-report")
240
241
242 @app.get("/")
243 def hello():
244 return {"message": "CoPilot - We Made It!"}
245
246
247 if __name__ == "__main__":
248 uvicorn.run(app, host=server_ip, port=5000)