Manual update agent client (#217)

* feat: Add optional fields for Velociraptor agent in Agents model

The code changes in `universal_models.py` modify the `Agents` model to include optional fields for the Velociraptor agent. The `velociraptor_id` and `velociraptor_last_seen` fields are now optional, allowing for cases where the Velociraptor agent is not present. Similarly, the `velociraptor_agent_version` field is also made optional. This change improves the flexibility of the model and accommodates scenarios where the Velociraptor agent may not be used.

Note: This commit message follows the established convention of using a prefix to indicate the type of change (`feat` for a new feature) and provides a clear and concise description of the changes made.

* agent rewrite

* feat: Add endpoint to update agent's Velociraptor ID

The code changes in `agents.py` add a new endpoint `/update` to update an agent's Velociraptor ID. This endpoint requires the `agent_id` and `velociraptor_id` as parameters and updates the corresponding agent's `velociraptor_id` field in the database. If the agent is not found, a 404 error is returned. This feature improves the functionality of the application by allowing users to easily update the Velociraptor ID of an agent.

Note: This commit message follows the established convention of using a prefix to indicate the type of change (`feat` for a new feature) and provides a clear and concise description of the changes made.

* feat: Refactor agent_sync to remove unnecessary session parameter

The code changes in `agent_sync.py` remove the unnecessary `session` parameter from the `sync_all_agents` function. Since the function is now using the `get_db_session` context manager, there is no need to pass the session as a parameter. This refactor simplifies the code and improves readability.

Note: This commit message follows the established convention of using a prefix to indicate the type of change (`feat` for a new feature) and provides a clear and concise description of the changes made.

* updated dependencies

* updated agent api

* added AgentVelociraptorIdForm

* precommit fixes

---------

Co-authored-by: Davide Di Modica <webmaster.ddm@gmail.com>

backend/alembic/versions/39c3aaec0084_new_agents_table_schema.py

@@ -0,0 +1,35 @@
+"""New agents table schema
+
+Revision ID: 39c3aaec0084
+Revises: ec63589cc24d
+Create Date: 2024-05-17 17:17:39.140779
+
+"""
+from typing import Sequence
+from typing import Union
+
+from sqlalchemy.dialects import mysql
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "39c3aaec0084"
+down_revision: Union[str, None] = "ec63589cc24d"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column("agents", "velociraptor_id", existing_type=mysql.VARCHAR(length=256), nullable=True)
+    op.alter_column("agents", "velociraptor_last_seen", existing_type=mysql.DATETIME(), nullable=True)
+    op.alter_column("agents", "velociraptor_agent_version", existing_type=mysql.VARCHAR(length=256), nullable=True)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column("agents", "velociraptor_agent_version", existing_type=mysql.VARCHAR(length=256), nullable=False)
+    op.alter_column("agents", "velociraptor_last_seen", existing_type=mysql.DATETIME(), nullable=False)
+    op.alter_column("agents", "velociraptor_id", existing_type=mysql.VARCHAR(length=256), nullable=False)
+    # ### end Alembic commands ###

backend/app/agents/routes/agents.py

@@ -18,7 +18,8 @@
 from app.agents.schema.agents import SyncedAgentsResponse
 from app.agents.services.status import get_outdated_agents_velociraptor
 from app.agents.services.status import get_outdated_agents_wazuh
-from app.agents.services.sync import sync_agents
+from app.agents.services.sync import sync_agents_velociraptor
+from app.agents.services.sync import sync_agents_wazuh
 from app.agents.velociraptor.services.agents import delete_agent_velociraptor
 from app.agents.wazuh.schema.agents import WazuhAgentScaPolicyResultsResponse
 from app.agents.wazuh.schema.agents import WazuhAgentScaResponse
@@ -223,10 +224,7 @@ async def get_agent_by_hostname(
         Security(AuthHandler().require_any_scope("admin", "analyst", "scheduler")),
     ],
 )
-async def sync_all_agents(
-    # backgroud_tasks: BackgroundTasks,
-    session: AsyncSession = Depends(get_db),
-) -> SyncedAgentsResponse:
+async def sync_all_agents() -> SyncedAgentsResponse:
     """
     Sync all agents from Wazuh Manager.
 
@@ -241,10 +239,10 @@ async def sync_all_agents(
     - SyncedAgentsResponse: The response model indicating the success of the sync operation.
 
     """
-    logger.info("Syncing agents from Wazuh Manager")
-    # backgroud_tasks.add_task(sync_agents, session)
+    logger.info("Syncing agents as part of scheduled job")
     loop = asyncio.get_event_loop()
-    loop.create_task(sync_agents(session=session))
+    await loop.create_task(sync_agents_wazuh())
+    await loop.create_task(sync_agents_velociraptor())
     return SyncedAgentsResponse(
         success=True,
         message="Agents synced started successfully",
@@ -474,7 +472,51 @@ async def get_outdated_velociraptor_agents(
     return await get_outdated_agents_velociraptor(session)
 
 
-# ! TODO: FINISH THIS
+@agents_router.put(
+    "/{agent_id}/update",
+    response_model=AgentModifyResponse,
+    description="Update agent",
+    dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
+)
+async def update_agent(
+    agent_id: str,
+    velociraptor_id: str,
+    session: AsyncSession = Depends(get_db),
+) -> AgentModifyResponse:
+    """
+    Updates an agent's velociraptor_id
+
+    Args:
+        agent_id (str): The ID of the agent to be updated.
+        velociraptor_id (str): The new velociraptor_id of the agent.
+        session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
+
+    Returns:
+        AgentModifyResponse: The response indicating the success or failure of the update.
+    """
+    logger.info(f"Updating agent {agent_id} with Velociraptor ID: {velociraptor_id}")
+    try:
+        result = await session.execute(select(Agents).filter(Agents.agent_id == agent_id))
+        agent = result.scalars().first()
+        if not agent:
+            raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
+        agent.velociraptor_id = velociraptor_id
+        await session.commit()
+        logger.info(f"Agent {agent_id} updated with Velociraptor ID: {velociraptor_id}")
+        return AgentModifyResponse(
+            success=True,
+            message=f"Agent {agent_id} updated with Velociraptor ID: {velociraptor_id}",
+        )
+    except Exception as e:
+        if not agent:
+            raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
+        logger.error(f"Failed to update agent {agent_id} with Velociraptor ID: {velociraptor_id}: {e}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to update agent {agent_id} with Velociraptor ID: {velociraptor_id}: {e}",
+        )
+
+
 @agents_router.delete(
     "/{agent_id}/delete",
     response_model=AgentModifyResponse,

backend/app/agents/schema/agents.py

@@ -18,6 +18,10 @@ class SyncedAgent(WazuhAgent, VelociraptorAgent):
     pass
 
 
+class SyncedWazuhAgent(WazuhAgent):
+    pass
+
+
 class SyncedAgentsResponse(BaseModel):
     # agents_added: List[SyncedAgent]
     success: bool