Agent checkins (#618)

* initial work on checkin data

* fix starkiller commit

* update sqlalchemy, use hybrid expressins for stale and lastseen

* fix tests

* gettinga bit more efficient

* fix tests post-merge

* fixes for test shutdown, insert checkins faster

* update performance test

* refactoring test fixtures

* refactor tests, stub aggregate endpoints

* get basic aggregate functionality

* fix formatting

* fix aggregate tests, disable aggregeate functions for sqlite'

* test timezone handling

* changelog

* working through cleaning up diff

* add a 'slow' marker for tests. update pytest to v7 minimum

* fix some todos, add others

* assert response time

* make 3.8 compatible

* calm the logging

* try to fix reset hanging issue

* adjust response times for slower machines

* give more breathing room to the slower ci machines

* fix 204 response

* fix database deletion

* upgrade sqlalchemy-utc to suppress all the cache_ok warnings

* bump timeout to 30 min temporarily

* remove a couple duplicate steps

* fix changelog

.github/CONTRIBUTING.md

@@ -44,4 +44,10 @@ The more information you provide in a Github issue the easier it will be for us
     3. run `ruff . --fix` (or `poetry run ruff . --fix`). 
 * The repo is also configured to use [pre-commit](https://pre-commit.com/) to automatically format code.
   * Once you have pre-commit installed, you can run `pre-commit install` to install the pre-commit hooks.
-  * Then pre-commit will execute black, isort, and ruff automatically before committing.
+  * Then pre-commit will execute black, isort, and ruff automatically before committing.
+
+## Tests
+
+Please write tests for your code! We use [pytest](https://docs.pytest.org/en/latest/) for testing. Tests are located in the `tests/` directory. To run the tests, run `pytest` from the root directory of the project.
+
+For tests that take >20-30 seconds, please add the `@pytest.mark.slow` decorator to the test function. This will allow us to skip the slow tests when running the tests, unless we explicitly want to run them with `pytest --runslow`.

.github/docker-compose.yml

@@ -13,9 +13,9 @@ services:
     entrypoint: /bin/bash
     platform: linux/amd64
     command: >
-      -c "DATABASE_USE=sqlite poetry run python -m pytest .
+      -c "DATABASE_USE=sqlite poetry run python -m pytest . --runslow
       && sed -i 's/localhost:3306/db:3306/g' empire/test/test_server_config.yaml
-      && DATABASE_USE=mysql poetry run python -m pytest ."
+      && DATABASE_USE=mysql poetry run python -m pytest . --runslow"
   db:
     image: mysql:8.0
     restart: always

.github/workflows/lint-and-test.yml

@@ -37,7 +37,7 @@ jobs:
     needs:
       - matrix-prep-config
       - lint
-    timeout-minutes: 15
+    timeout-minutes: 30
     runs-on: ubuntu-latest
     name: Test Python ${{ matrix.python-version }}
     strategy:
@@ -76,11 +76,11 @@ jobs:
           poetry install
       - name: Run test suite - mysql
         run: |
-          DATABASE_USE=mysql poetry run pytest . -v
+          DATABASE_USE=mysql poetry run pytest . -v --runslow
       - name: Run test suite - sqlite
         if: ${{ startsWith(github.head_ref, 'release/') || contains(github.event.pull_request.labels.*.name, 'test-sqlite') }}
         run: |
-          DATABASE_USE=sqlite poetry run pytest . -v
+          DATABASE_USE=sqlite poetry run pytest . -v --runslow
 
   test_image:
     # To save CI time, only run these tests on the release PRs

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+-   Break out agent checkins to a new table (@Vinnybod)
+    -   New checkins endpoint to get them as a list
+    -   New checkins aggregate endpoint to get aggregated checkin data
+    -   Aggregate endpoint not supported with SQLite
+-   Add a warning message about using SQLite
 -   Added LinPEAS to Python modules (@Cx01N)
 -   Added python obfusscation using python-obfuscator (@Cx01N)
 -   Added IronPython SMB Agents/Listener (@Cx01N)

conftest.py

@@ -0,0 +1,21 @@
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--runslow", action="store_true", default=False, help="run slow tests"
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "slow: mark test as slow to run")
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--runslow"):
+        # --runslow given in cli: do not skip slow tests
+        return
+    skip_slow = pytest.mark.skip(reason="need --runslow option to run")
+    for item in items:
+        if "slow" in item.keywords:
+            item.add_marker(skip_slow)

empire/server/api/v2/agent/agent_api.py

@@ -1,16 +1,30 @@
-from fastapi import Depends, HTTPException
+import math
+from datetime import datetime
+from typing import List, Optional
+
+from fastapi import Depends, HTTPException, Query
 from sqlalchemy.orm import Session
 
 from empire.server.api.api_router import APIRouter
 from empire.server.api.jwt_auth import get_current_active_user
 from empire.server.api.v2.agent.agent_dto import (
     Agent,
+    AgentCheckIns,
+    AgentCheckInsAggregate,
     Agents,
     AgentUpdateRequest,
+    AggregateBucket,
     domain_to_dto_agent,
+    domain_to_dto_agent_checkin,
+    domain_to_dto_agent_checkin_agg,
 )
 from empire.server.api.v2.shared_dependencies import get_db
-from empire.server.api.v2.shared_dto import BadRequestResponse, NotFoundResponse
+from empire.server.api.v2.shared_dto import (
+    BadRequestResponse,
+    NotFoundResponse,
+    OrderDirection,
+)
+from empire.server.core.config import empire_config
 from empire.server.core.db import models
 from empire.server.server import main
 
@@ -36,6 +50,57 @@ async def get_agent(uid: str, db: Session = Depends(get_db)):
     raise HTTPException(404, f"Agent not found for id {uid}")
 
 
+@router.get("/checkins", response_model=AgentCheckIns)
+def read_agent_checkins_all(
+    db: Session = Depends(get_db),
+    agents: List[str] = Query(None),
+    limit: int = 1000,
+    page: int = 1,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    order_direction: OrderDirection = OrderDirection.desc,
+):
+    checkins, total = agent_service.get_agent_checkins(
+        db, agents, limit, (page - 1) * limit, start_date, end_date, order_direction
+    )
+    checkins = list(map(lambda x: domain_to_dto_agent_checkin(x), checkins))
+
+    return AgentCheckIns(
+        records=checkins,
+        page=page,
+        total_pages=math.ceil(total / limit),
+        limit=limit,
+        total=total,
+    )
+
+
+@router.get("/checkins/aggregate", response_model=AgentCheckInsAggregate)
+def read_agent_checkins_aggregate(
+    db: Session = Depends(get_db),
+    agents: List[str] = Query(None),
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    bucket_size: Optional[AggregateBucket] = AggregateBucket.day,
+):
+    if empire_config.database.use == "sqlite":
+        raise HTTPException(
+            400,
+            "Aggregate checkins not supported with sqlite. Please use MySQL.",
+        )
+
+    checkins = agent_service.get_agent_checkins_aggregate(
+        db, agents, start_date, end_date, bucket_size
+    )
+    checkins = list(map(lambda x: domain_to_dto_agent_checkin_agg(x), checkins))
+
+    return AgentCheckInsAggregate(
+        records=checkins,
+        start_date=start_date,
+        end_date=end_date,
+        bucket_size=bucket_size,
+    )
+
+
 @router.get("/{uid}", response_model=Agent)
 async def read_agent(uid: str, db_agent: models.Agent = Depends(get_agent)):
     return domain_to_dto_agent(db_agent)
@@ -70,3 +135,33 @@ async def update_agent(
         raise HTTPException(status_code=400, detail=err)
 
     return domain_to_dto_agent(resp)
+
+
+@router.get("/{uid}/checkins", response_model=AgentCheckIns)
+def read_agent_checkins(
+    db: Session = Depends(get_db),
+    db_agent: models.Agent = Depends(get_agent),
+    limit: int = -1,
+    page: int = 1,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    order_direction: OrderDirection = OrderDirection.desc,
+):
+    checkins, total = agent_service.get_agent_checkins(
+        db,
+        [db_agent.session_id],
+        limit,
+        (page - 1) * limit,
+        start_date,
+        end_date,
+        order_direction,
+    )
+    checkins = list(map(lambda x: domain_to_dto_agent_checkin(x), checkins))
+
+    return AgentCheckIns(
+        records=checkins,
+        page=page,
+        total_pages=math.ceil(total / limit),
+        limit=limit,
+        total=total,
+    )

empire/server/api/v2/agent/agent_dto.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from enum import Enum
 from typing import Dict, List, Optional
 
 from pydantic import BaseModel
@@ -31,7 +32,7 @@ def domain_to_dto_agent(agent: models.Agent):
         process_name=agent.process_name,
         os_details=agent.os_details,
         nonce=agent.nonce,
-        checkin_time=agent.checkin_time,
+        checkin_time=agent.firstseen_time,
         lastseen_time=agent.lastseen_time,
         parent=agent.parent,
         children=agent.children,
@@ -63,6 +64,19 @@ def to_proxy_dto(proxies):
     return {}
 
 
+def domain_to_dto_agent_checkin(agent_checkin: models.AgentCheckIn):
+    return AgentCheckIn(
+        agent_id=agent_checkin.agent_id,
+        checkin_time=agent_checkin.checkin_time,
+    )
+
+
+def domain_to_dto_agent_checkin_agg(agent_checkin_agg):
+    return AgentCheckInAggregate(
+        count=agent_checkin_agg["count"], checkin_time=agent_checkin_agg["checkin_time"]
+    )
+
+
 class Agent(BaseModel):
     session_id: str
     name: str
@@ -103,6 +117,38 @@ class Agents(BaseModel):
     records: List[Agent]
 
 
+class AgentCheckIn(BaseModel):
+    agent_id: str
+    checkin_time: datetime
+
+
+class AgentCheckIns(BaseModel):
+    records: List[AgentCheckIn]
+    limit: int
+    page: int
+    total_pages: int
+    total: int
+
+
+class AgentCheckInAggregate(BaseModel):
+    count: int
+    checkin_time: datetime  # will be truncated depending on the group_by
+
+
+class AgentCheckInsAggregate(BaseModel):
+    records: List[AgentCheckInAggregate]
+    start_date: Optional[datetime]
+    end_date: Optional[datetime]
+    bucket_size: str
+
+
+class AggregateBucket(str, Enum):
+    second = "second"
+    minute = "minute"
+    hour = "hour"
+    day = "day"
+
+
 class AgentUpdateRequest(BaseModel):
     name: str
     notes: Optional[str]

empire/server/api/v2/module/module_api.py

@@ -1,7 +1,6 @@
 import logging
-from datetime import datetime
 
-from fastapi import Depends, HTTPException
+from fastapi import Depends, HTTPException, Response
 from sqlalchemy.orm import Session
 
 from empire.server.api.api_router import APIRouter
@@ -49,15 +48,12 @@ async def get_module(uid: str):
     # response_model=Modules,
 )
 async def read_modules():
-    log.info(f"Request Received {datetime.utcnow()}")
     modules = list(
         map(
             lambda x: domain_to_dto_module(x[1], x[0]), module_service.get_all().items()
         )
     )
 
-    log.info(f"Done Converting Objects {datetime.utcnow()}")
-
     return {"records": modules}
 
 
@@ -88,7 +84,7 @@ async def update_module(
     return domain_to_dto_module(module, uid)
 
 
-@router.put("/bulk/enable", status_code=204)
+@router.put("/bulk/enable", status_code=204, response_class=Response)
 async def update_bulk_enable(
     module_req: ModuleBulkUpdateRequest, db: Session = Depends(get_db)
 ):

empire/server/common/agents.py

@@ -46,7 +46,7 @@
 from pathlib import Path
 from typing import Dict
 
-from sqlalchemy import and_, or_, update
+from sqlalchemy import and_, or_
 from sqlalchemy.orm import Session
 from zlib_wrapper import decompress
 
@@ -185,6 +185,7 @@ def add_agent(
         )
 
         db.add(agent)
+        db.add(models.AgentCheckIn(agent_id=sessionID))
         db.flush()
 
         message = f"New agent {sessionID} checked in"
@@ -688,19 +689,7 @@ def update_agent_lastseen_db(self, session_id, db: Session):
         """
         Update the agent's last seen timestamp in the database.
         """
-        warnings.warn(
-            "This has been deprecated and may be removed."
-            "Use agent_service.get_by_id() or agent_service.get_by_name() instead.",
-            DeprecationWarning,
-        )
-        db.execute(
-            update(models.Agent).where(
-                or_(
-                    models.Agent.session_id == session_id,
-                    models.Agent.name == session_id,
-                )
-            )
-        )
+        db.add(models.AgentCheckIn(agent_id=session_id))
 
     def set_autoruns_db(self, task_command, module_data):
         """