Skip to content

Commit

Permalink
Tags (#649)
Browse files Browse the repository at this point in the history
* initial tags

* add tags endpoint for listeners with tests

* update tag

* expand to agents

* fixing some things

* fix agent tags, fix listener api tests

* abstract tag tests

* abstract tags endpoints

* add plugin task tags

* add tags to credentials

* add tags to downloads

* add events for tags

* refactor. working on list endpoint

* more cleanup, basic list endpoint

* add validation

* add ordering to tags

* remove starkiller submodule

* remove todo

* more tags tests. add tags to download filters

* fix stager test cleanup

* add label to tag dto

* fix user filters

* fix reset test
  • Loading branch information
vinnybod authored Jul 17, 2023
1 parent 98f791c commit af5f92e
Show file tree
Hide file tree
Showing 29 changed files with 1,389 additions and 146 deletions.
2 changes: 2 additions & 0 deletions empire/server/api/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ def initialize(secure: bool = False, ip: str = "0.0.0.0", port: int = 1337):
from empire.server.api.v2.plugin import plugin_api, plugin_task_api
from empire.server.api.v2.profile import profile_api
from empire.server.api.v2.stager import stager_api, stager_template_api
from empire.server.api.v2.tag import tag_api
from empire.server.api.v2.user import user_api
from empire.server.server import main

Expand Down Expand Up @@ -97,6 +98,7 @@ def shutdown_event():
v2App.include_router(meta_api.router)
v2App.include_router(plugin_task_api.router)
v2App.include_router(plugin_api.router)
v2App.include_router(tag_api.router)

v2App.add_middleware(
EmpireCORSMiddleware,
Expand Down
4 changes: 4 additions & 0 deletions empire/server/api/v2/agent/agent_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
NotFoundResponse,
OrderDirection,
)
from empire.server.api.v2.tag import tag_api
from empire.server.core.config import empire_config
from empire.server.core.db import models
from empire.server.server import main
Expand All @@ -50,6 +51,9 @@ async def get_agent(uid: str, db: Session = Depends(get_db)):
raise HTTPException(404, f"Agent not found for id {uid}")


tag_api.add_endpoints_to_taggable(router, "/{uid}/tags", get_agent)


@router.get("/checkins", response_model=AgentCheckIns)
def read_agent_checkins_all(
db: Session = Depends(get_db),
Expand Down
3 changes: 3 additions & 0 deletions empire/server/api/v2/agent/agent_dto.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from pydantic import BaseModel

from empire.server.api.v2.shared_dto import PROXY_ID
from empire.server.api.v2.tag.tag_dto import Tag, domain_to_dto_tag
from empire.server.core.db import models


Expand Down Expand Up @@ -48,6 +49,7 @@ def domain_to_dto_agent(agent: models.Agent):
archived=agent.archived,
# Could make this a typed class later to match the schema
proxies=to_proxy_dto(agent.proxies),
tags=list(map(lambda x: domain_to_dto_tag(x), agent.tags)),
)


Expand Down Expand Up @@ -111,6 +113,7 @@ class Agent(BaseModel):
archived: bool
stale: bool
proxies: Optional[Dict]
tags: List[Tag]


class Agents(BaseModel):
Expand Down
9 changes: 9 additions & 0 deletions empire/server/api/v2/agent/agent_task_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@
NotFoundResponse,
OrderDirection,
)
from empire.server.api.v2.tag import tag_api
from empire.server.api.v2.tag.tag_dto import TagStr
from empire.server.core.agent_service import AgentService
from empire.server.core.agent_task_service import AgentTaskService
from empire.server.core.db import models
Expand Down Expand Up @@ -83,6 +85,9 @@ async def get_task(
)


tag_api.add_endpoints_to_taggable(router, "/{agent_id}/tasks/{uid}/tags", get_task)


@router.get("/tasks", response_model=AgentTasks)
async def read_tasks_all_agents(
limit: int = -1,
Expand All @@ -96,13 +101,15 @@ async def read_tasks_all_agents(
status: Optional[AgentTaskStatus] = None,
agents: Optional[List[str]] = Query(None),
users: Optional[List[int]] = Query(None),
tags: Optional[List[TagStr]] = Query(None),
query: Optional[str] = None,
db: Session = Depends(get_db),
):
tasks, total = agent_task_service.get_tasks(
db,
agents=agents,
users=users,
tags=tags,
limit=limit,
offset=(page - 1) * limit,
include_full_input=include_full_input,
Expand Down Expand Up @@ -145,6 +152,7 @@ async def read_tasks(
order_direction: OrderDirection = OrderDirection.desc,
status: Optional[AgentTaskStatus] = None,
users: Optional[List[int]] = Query(None),
tags: Optional[List[TagStr]] = Query(None),
db: Session = Depends(get_db),
db_agent: models.Agent = Depends(get_agent),
query: Optional[str] = None,
Expand All @@ -153,6 +161,7 @@ async def read_tasks(
db,
agents=[db_agent.session_id],
users=users,
tags=tags,
limit=limit,
offset=(page - 1) * limit,
include_full_input=include_full_input,
Expand Down
3 changes: 3 additions & 0 deletions empire/server/api/v2/agent/agent_task_dto.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
DownloadDescription,
domain_to_dto_download_description,
)
from empire.server.api.v2.tag.tag_dto import Tag, domain_to_dto_tag
from empire.server.core.db import models


Expand Down Expand Up @@ -41,6 +42,7 @@ def domain_to_dto_task(
status=task.status,
created_at=task.created_at,
updated_at=task.updated_at,
tags=list(map(lambda x: domain_to_dto_tag(x), task.tags)),
)


Expand All @@ -59,6 +61,7 @@ class AgentTask(BaseModel):
status: models.AgentTaskStatus
created_at: datetime
updated_at: datetime
tags: List[Tag]


class AgentTasks(BaseModel):
Expand Down
4 changes: 4 additions & 0 deletions empire/server/api/v2/credential/credential_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
)
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.tag import tag_api
from empire.server.core.db import models
from empire.server.server import main

Expand All @@ -41,6 +42,9 @@ async def get_credential(uid: int, db: Session = Depends(get_db)):
raise HTTPException(404, f"Credential not found for id {uid}")


tag_api.add_endpoints_to_taggable(router, "/{uid}/tags", get_credential)


@router.get("/{uid}", response_model=Credential)
async def read_credential(
uid: int, db_credential: models.Credential = Depends(get_credential)
Expand Down
4 changes: 4 additions & 0 deletions empire/server/api/v2/credential/credential_dto.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@

from pydantic import BaseModel

from empire.server.api.v2.tag.tag_dto import Tag, domain_to_dto_tag


def domain_to_dto_credential(credential):
return Credential(
Expand All @@ -17,6 +19,7 @@ def domain_to_dto_credential(credential):
notes=credential.notes,
created_at=credential.created_at,
updated_at=credential.updated_at,
tags=list(map(lambda x: domain_to_dto_tag(x), credential.tags)),
)


Expand All @@ -32,6 +35,7 @@ class Credential(BaseModel):
notes: Optional[str]
created_at: datetime
updated_at: datetime
tags: List[Tag]


class Credentials(BaseModel):
Expand Down
8 changes: 8 additions & 0 deletions empire/server/api/v2/download/download_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
NotFoundResponse,
OrderDirection,
)
from empire.server.api.v2.tag import tag_api
from empire.server.api.v2.tag.tag_dto import TagStr
from empire.server.core.db import models
from empire.server.server import main

Expand Down Expand Up @@ -59,8 +61,12 @@ async def download_download(
return FileResponse(db_download.location, filename=filename)


tag_api.add_endpoints_to_taggable(router, "/{uid}/tags", get_download)


@router.get(
"/{uid}",
response_model=Download,
)
async def read_download(
uid: int,
Expand All @@ -79,10 +85,12 @@ async def read_downloads(
order_by: DownloadOrderOptions = DownloadOrderOptions.updated_at,
query: Optional[str] = None,
sources: Optional[List[DownloadSourceFilter]] = Query(None),
tags: Optional[List[TagStr]] = Query(None),
):
downloads, total = download_service.get_all(
db=db,
download_types=sources,
tags=tags,
q=query,
limit=limit,
offset=(page - 1) * limit,
Expand Down
4 changes: 4 additions & 0 deletions empire/server/api/v2/download/download_dto.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

from pydantic import BaseModel

from empire.server.api.v2.tag.tag_dto import Tag, domain_to_dto_tag


def removeprefix(value: str, prefix: str) -> str:
if value.startswith(prefix):
Expand All @@ -21,6 +23,7 @@ def domain_to_dto_download(download):
size=download.size,
created_at=download.created_at,
updated_at=download.updated_at,
tags=list(map(lambda x: domain_to_dto_tag(x), download.tags)),
)


Expand All @@ -46,6 +49,7 @@ class Download(BaseModel):
size: int
created_at: datetime
updated_at: datetime
tags: List[Tag]


class Downloads(BaseModel):
Expand Down
4 changes: 4 additions & 0 deletions empire/server/api/v2/listener/listener_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
)
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.tag import tag_api
from empire.server.core.db import models
from empire.server.server import main

Expand All @@ -39,6 +40,9 @@ async def get_listener(uid: int, db: Session = Depends(get_db)):
raise HTTPException(404, f"Listener not found for id {uid}")


tag_api.add_endpoints_to_taggable(router, "/{uid}/tags", get_listener)


@router.get("/{uid}", response_model=Listener)
async def read_listener(uid: int, db_listener: models.Listener = Depends(get_listener)):
return domain_to_dto_listener(db_listener)
Expand Down
3 changes: 3 additions & 0 deletions empire/server/api/v2/listener/listener_dto.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from pydantic import BaseModel

from empire.server.api.v2.shared_dto import Author, CustomOptionSchema, to_value_type
from empire.server.api.v2.tag.tag_dto import Tag, domain_to_dto_tag


def domain_to_dto_template(listener, uid: str):
Expand Down Expand Up @@ -59,6 +60,7 @@ def domain_to_dto_listener(listener):
enabled=listener.enabled,
options=options,
created_at=listener.created_at,
tags=list(map(lambda x: domain_to_dto_tag(x), listener.tags)),
)


Expand Down Expand Up @@ -249,6 +251,7 @@ class Listener(BaseModel):
template: str
options: Dict[str, str]
created_at: datetime
tags: List[Tag]


class Listeners(BaseModel):
Expand Down
9 changes: 9 additions & 0 deletions empire/server/api/v2/plugin/plugin_task_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
NotFoundResponse,
OrderDirection,
)
from empire.server.api.v2.tag import tag_api
from empire.server.api.v2.tag.tag_dto import TagStr
from empire.server.core.db import models
from empire.server.core.db.models import PluginTaskStatus
from empire.server.core.download_service import DownloadService
Expand Down Expand Up @@ -59,6 +61,9 @@ async def get_task(uid: int, db: Session = Depends(get_db), plugin=Depends(get_p
)


tag_api.add_endpoints_to_taggable(router, "/{plugin_id}/tasks/{uid}/tags", get_task)


@router.get("/tasks", response_model=PluginTasks)
async def read_tasks_all_plugins(
limit: int = -1,
Expand All @@ -71,13 +76,15 @@ async def read_tasks_all_plugins(
status: Optional[PluginTaskStatus] = None,
plugins: Optional[List[str]] = Query(None),
users: Optional[List[int]] = Query(None),
tags: Optional[List[TagStr]] = Query(None),
query: Optional[str] = None,
db: Session = Depends(get_db),
):
tasks, total = plugin_service.get_tasks(
db,
plugins=plugins,
users=users,
tags=tags,
limit=limit,
offset=(page - 1) * limit,
include_full_input=include_full_input,
Expand Down Expand Up @@ -116,6 +123,7 @@ async def read_tasks(
order_direction: OrderDirection = OrderDirection.desc,
status: Optional[PluginTaskStatus] = None,
users: Optional[List[int]] = Query(None),
tags: Optional[List[TagStr]] = Query(None),
db: Session = Depends(get_db),
plugin=Depends(get_plugin),
query: Optional[str] = None,
Expand All @@ -124,6 +132,7 @@ async def read_tasks(
db,
plugins=[plugin.info["Name"]],
users=users,
tags=tags,
limit=limit,
offset=(page - 1) * limit,
include_full_input=include_full_input,
Expand Down
3 changes: 3 additions & 0 deletions empire/server/api/v2/plugin/plugin_task_dto.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
DownloadDescription,
domain_to_dto_download_description,
)
from empire.server.api.v2.tag.tag_dto import Tag, domain_to_dto_tag
from empire.server.core.db import models


Expand Down Expand Up @@ -37,6 +38,7 @@ def domain_to_dto_plugin_task(
status=task.status,
created_at=task.created_at,
updated_at=task.updated_at,
tags=list(map(lambda x: domain_to_dto_tag(x), task.tags)),
)


Expand All @@ -52,6 +54,7 @@ class PluginTask(BaseModel):
status: Optional[models.PluginTaskStatus]
created_at: datetime
updated_at: datetime
tags: List[Tag]


class PluginTasks(BaseModel):
Expand Down
Empty file.
Loading

0 comments on commit af5f92e

Please sign in to comment.