Bypass language, stale processes, keyword length (#452)

* add minimum requirements for keyword dto

* add language to bypass endpoints. Update ps hook to mark processes stale. add requirements to keyword dto

* add stale process to endpoint, fix int comparison

* delete hostprocesses after hook test

empire/server/api/v2/bypass/bypass_dto.py

@@ -11,6 +11,7 @@ def domain_to_dto_bypass(bypass):
         id=bypass.id,
         name=bypass.name,
         authors=bypass.authors or [],
+        language=bypass.language,
         code=bypass.code,
         created_at=bypass.created_at,
         updated_at=bypass.updated_at,
@@ -21,6 +22,7 @@ class Bypass(BaseModel):
     id: int
     name: str
     authors: List[Author]
+    language: str
     code: str
     created_at: datetime
     updated_at: datetime
@@ -32,9 +34,11 @@ class Bypasses(BaseModel):
 
 class BypassUpdateRequest(BaseModel):
     name: str
+    language: str
     code: str
 
 
 class BypassPostRequest(BaseModel):
     name: str
+    language: str
     code: str

empire/server/api/v2/host/process_dto.py

@@ -17,6 +17,7 @@ def domain_to_dto_process(process: models.HostProcess):
         host_id=process.host_id,
         architecture=process.architecture,
         user=process.user,
+        stale=process.stale,
         agent_id=agent_id,
     )
 
@@ -27,6 +28,7 @@ class Process(BaseModel):
     host_id: int
     architecture: Optional[str]
     user: Optional[str]
+    stale: bool
     agent_id: Optional[str]
 
 

empire/server/api/v2/obfuscation/obfuscation_dto.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import List
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from empire.server.core.db import models
 
@@ -22,13 +22,13 @@ class Keywords(BaseModel):
 
 
 class KeywordUpdateRequest(BaseModel):
-    keyword: str
-    replacement: str
+    keyword: str = Field(min_length=3)
+    replacement: str = Field(min_length=3)
 
 
 class KeywordPostRequest(BaseModel):
-    keyword: str
-    replacement: str
+    keyword: str = Field(min_length=3)
+    replacement: str = Field(min_length=3)
 
 
 def domain_to_dto_obfuscation_config(obf_conf: models.ObfuscationConfig):

empire/server/core/bypass_service.py

@@ -80,7 +80,9 @@ def create_bypass(self, db: Session, bypass_req):
         if self.get_by_name(db, bypass_req.name):
             return None, f"Bypass with name {bypass_req.name} already exists."
 
-        bypass = models.Bypass(name=bypass_req.name, code=bypass_req.code)
+        bypass = models.Bypass(
+            name=bypass_req.name, code=bypass_req.code, language=bypass_req.language
+        )
 
         db.add(bypass)
         db.flush()
@@ -95,6 +97,7 @@ def update_bypass(self, db: Session, db_bypass: models.Bypass, bypass_req):
                 return None, f"Bypass with name {bypass_req.name} already exists."
 
         db_bypass.code = bypass_req.code
+        db_bypass.language = bypass_req.language
 
         db.flush()
 

empire/server/core/db/models.py

@@ -173,6 +173,7 @@ class HostProcess(Base):
     process_name = Column(Text)
     architecture = Column(String(255))
     user = Column(String(255))
+    stale = Column(Boolean, default=False)
     agent = relationship(
         Agent,
         lazy="joined",

empire/server/core/hooks_internal.py

@@ -55,7 +55,7 @@ def ps_hook(db: Session, tasking: models.Tasking):
         arch = process.get("Arch")
         user = process.get("UserName")
         if process_id:
-            # and process_id.isnumeric():
+            # new process
             if int(process_id) not in existing_processes:
                 db.add(
                     models.HostProcess(
@@ -66,6 +66,7 @@ def ps_hook(db: Session, tasking: models.Tasking):
                         user=user,
                     )
                 )
+            # update existing process
             elif int(process_id) in existing_processes:
                 db_process: models.HostProcess = (
                     db.query(models.HostProcess)
@@ -80,6 +81,22 @@ def ps_hook(db: Session, tasking: models.Tasking):
                 if not db_process.agent:
                     db_process.architecture = arch
                     db_process.process_name = process_name
+                    db_process.user = user
+
+    for process in existing_processes:
+        # mark processes that are no longer running stale
+        if process not in list(map(lambda p: int(p.get("PID")), output)):
+            db_process: models.HostProcess = (
+                db.query(models.HostProcess)
+                .filter(
+                    and_(
+                        models.HostProcess.host_id == tasking.agent.host_id,
+                        models.HostProcess.process_id == process,
+                    )
+                )
+                .first()
+            )
+            db_process.stale = True
 
 
 def ps_filter(db: Session, tasking: models.Tasking):

empire/test/test_bypass_api.py

@@ -24,7 +24,7 @@ def test_create_bypass_name_conflict(client, admin_auth_header):
     response = client.post(
         "/api/v2/bypasses/",
         headers=admin_auth_header,
-        json={"name": "mattifestation", "code": "x=0;"},
+        json={"name": "mattifestation", "code": "x=0;", "language": "powershell"},
     )
 
     assert response.status_code == 400
@@ -37,7 +37,7 @@ def test_create_bypass(client, admin_auth_header):
     response = client.post(
         "/api/v2/bypasses/",
         headers=admin_auth_header,
-        json={"name": "Test Bypass", "code": "x=0;"},
+        json={"name": "Test Bypass", "code": "x=0;", "language": "powershell"},
     )
 
     assert response.status_code == 201
@@ -63,7 +63,7 @@ def test_update_bypass_name_conflict(client, admin_auth_header):
     response = client.put(
         f"/api/v2/bypasses/5",
         headers=admin_auth_header,
-        json={"name": bypass_1_name, "code": "x=0;"},
+        json={"name": bypass_1_name, "code": "x=0;", "language": "powershell"},
     )
 
     assert response.status_code == 400
@@ -76,7 +76,7 @@ def test_update_bypass(client, admin_auth_header):
     response = client.put(
         "/api/v2/bypasses/1",
         headers=admin_auth_header,
-        json={"name": "Updated Bypass", "code": "x=1;"},
+        json={"name": "Updated Bypass", "code": "x=1;", "language": "powershell"},
     )
 
     assert response.json()["name"] == "Updated Bypass"

empire/test/test_hooks_internal.py

@@ -0,0 +1,122 @@
+import json
+
+import pytest
+
+from empire.server.core.db.models import TaskingStatus
+from empire.server.core.hooks import hooks
+
+
+@pytest.fixture(scope="module", autouse=True)
+def agent(db, models, main):
+    hosts = db.query(models.Host).all()
+    if len(hosts) == 0:
+        host = models.Host(name="default_host", internal_ip="127.0.0.1")
+        db.add(host)
+    else:
+        host = hosts[0]
+
+    name = f'agent_{__name__.split(".")[-1]}'
+
+    agent = db.query(models.Agent).filter(models.Agent.session_id == name).first()
+    if not agent:
+        agent = models.Agent(
+            name=name,
+            session_id=name,
+            delay=1,
+            jitter=0.1,
+            external_ip="1.1.1.1",
+            session_key="qwerty",
+            nonce="nonce",
+            profile="profile",
+            kill_date="killDate",
+            working_hours="workingHours",
+            lost_limit=60,
+            listener="http",
+            language="powershell",
+            language_version="5",
+            high_integrity=True,
+            process_name="abc",
+            process_id=123,
+            hostname=host.name,
+            host_id=host.id,
+            archived=False,
+        )
+        db.add(agent)
+    else:
+        agent.archived = False
+
+    db.flush()
+
+    main.agents.agents[name] = {
+        "sessionKey": agent.session_key,
+        "functions": agent.functions,
+    }
+
+    yield agent
+
+    db.query(models.HostProcess).delete()
+    db.delete(agent)
+    db.delete(host)
+    db.commit()
+
+
+def test_ps_hook(client, db, models, agent):
+    existing_processes = [
+        models.HostProcess(
+            host_id=agent.host_id,
+            process_id=1,
+            process_name="should_be_stale",
+            architecture="x86",
+            user="test_user",
+        ),
+        models.HostProcess(
+            host_id=agent.host_id,
+            process_id=2,
+            process_name="should_be_updated",
+            architecture="x86",
+            user="test_user",
+        ),
+        models.HostProcess(
+            host_id=agent.host_id,
+            process_id=3,
+            process_name="should_be_same",
+            architecture="x86",
+            user="test_user",
+        ),
+    ]
+    db.add_all(existing_processes)
+
+    output = json.dumps(
+        [
+            {
+                "CMD": "has_been_updated",
+                "PID": 2,
+                "Arch": "x86_64",
+                "UserName": "test_user",
+            },
+            {"CMD": "should_be_same", "PID": 3, "Arch": "x86", "UserName": "test_user"},
+            {"CMD": "should_be_new", "PID": 4, "Arch": "x86", "UserName": "test_user"},
+        ]
+    )
+    tasking = models.Tasking(
+        id=1,
+        agent_id=agent.session_id,
+        agent=agent,
+        input="ps",
+        status=TaskingStatus.pulled,
+        output=output,
+        original_output=output,
+    )
+    hooks.run_hooks(hooks.BEFORE_TASKING_RESULT_HOOK, db, tasking)
+    db.flush()
+    processes = db.query(models.HostProcess).all()
+
+    assert len(processes) == 4
+    assert processes[0].process_name == "should_be_stale"
+    assert processes[0].stale is True
+    assert processes[1].process_name == "has_been_updated"
+    assert processes[1].stale is False
+    assert processes[2].process_name == "should_be_same"
+    assert processes[2].stale is False
+    assert processes[3].process_name == "should_be_new"
+    assert processes[3].stale is False

empire/test/test_obfuscation_api.py

@@ -53,6 +53,20 @@ def test_create_keyword_name_conflict(client, admin_auth_header):
     )
 
 
+def test_create_keyword_validate_length(client, admin_auth_header):
+    response = client.post(
+        "/api/v2/obfuscation/keywords/",
+        headers=admin_auth_header,
+        json={"keyword": "a", "replacement": "b"},
+    )
+
+    assert response.status_code == 422
+    assert (
+        response.json()["detail"][0]["msg"]
+        == "ensure this value has at least 3 characters"
+    )
+
+
 def test_create_keyword(client, admin_auth_header):
     response = client.post(
         "/api/v2/obfuscation/keywords/",