From cb2e4a00763b4e638bf13d15a3374de872753b51 Mon Sep 17 00:00:00 2001
From: Taylor <taylor.walton@socfortress.co>
Date: Sat, 2 Mar 2024 13:00:37 -0600
Subject: [PATCH] Add async support for agent synchronization. Remove the
 background task and run every 15 minutes

---
 backend/app/agents/routes/agents.py           |  7 +++++--
 backend/app/schedulers/scheduler.py           |  1 +
 backend/app/schedulers/services/agent_sync.py | 21 +++++--------------
 3 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/backend/app/agents/routes/agents.py b/backend/app/agents/routes/agents.py
index 6cf12022..cafeddf4 100644
--- a/backend/app/agents/routes/agents.py
+++ b/backend/app/agents/routes/agents.py
@@ -3,6 +3,7 @@
 from fastapi import Depends
 from fastapi import HTTPException
 from fastapi import Security
+import asyncio
 from loguru import logger
 from sqlalchemy import delete
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -218,7 +219,7 @@ async def get_agent_by_hostname(
     ],
 )
 async def sync_all_agents(
-    backgroud_tasks: BackgroundTasks,
+    #backgroud_tasks: BackgroundTasks,
     session: AsyncSession = Depends(get_db),
 ) -> SyncedAgentsResponse:
     """
@@ -236,7 +237,9 @@ async def sync_all_agents(
 
     """
     logger.info("Syncing agents from Wazuh Manager")
-    backgroud_tasks.add_task(sync_agents, session)
+    #backgroud_tasks.add_task(sync_agents, session)
+    loop = asyncio.get_event_loop()
+    loop.create_task(sync_agents(session))
     return SyncedAgentsResponse(
         success=True,
         message="Agents synced started successfully",
diff --git a/backend/app/schedulers/scheduler.py b/backend/app/schedulers/scheduler.py
index 34969e90..471ae8ca 100644
--- a/backend/app/schedulers/scheduler.py
+++ b/backend/app/schedulers/scheduler.py
@@ -85,6 +85,7 @@ def schedule_enabled_jobs(scheduler):
                     id=job_metadata.job_id,
                     replace_existing=True,
                 )
+                logger.info(f"Scheduled job: {job_metadata.job_id}")
             except ValueError as e:
                 logger.error(f"Error scheduling job: {e}")
 
diff --git a/backend/app/schedulers/services/agent_sync.py b/backend/app/schedulers/services/agent_sync.py
index f7cf5a21..d5f07208 100644
--- a/backend/app/schedulers/services/agent_sync.py
+++ b/backend/app/schedulers/services/agent_sync.py
@@ -6,12 +6,14 @@
 
 from app.db.db_session import get_sync_db_session
 from app.schedulers.models.scheduler import JobMetadata
+from app.db.db_session import get_db_session
 from app.schedulers.utils.universal import scheduler_login
+from app.agents.routes.agents import sync_all_agents
 
 load_dotenv()
 
 
-def agent_sync():
+async def agent_sync():
     """
     Synchronizes agents by sending a request to the server and updating the job metadata.
 
@@ -21,21 +23,8 @@ def agent_sync():
     If the token retrieval fails, it prints a failure message. If the job metadata for
     'agent_sync' does not exist, it prints a message indicating the absence of the metadata.
     """
-    # Get the scheduler auth token
-    headers = scheduler_login()
-
-    # Check if the token was successfully retrieved
-    if headers:
-        # Your actual task
-        response = requests.post(
-            f"http://{os.getenv('SERVER_IP')}:5000/agents/sync",
-            headers=headers,
-        )
-
-        # Process the response here if needed
-        print(response.json())
-    else:
-        print("Failed to retrieve token")
+    async with get_db_session() as session:
+        await sync_all_agents(session=session)
 
     # Use get_sync_db_session to create and manage a synchronous session
     with get_sync_db_session() as session: