diff --git a/api/crud.py b/api/crud.py index 332d5a5..724a08b 100644 --- a/api/crud.py +++ b/api/crud.py @@ -7,9 +7,13 @@ def get_items(db: Session, skip: int = 0, limit: int = 100): return db.query(models.Item).offset(skip).limit(limit).all() -def get_item(db: Session, item_id: int, query: str=""): +def get_item(db: Session, item_id: int, query: str = ""): if query: - return db.query(models.Item).filter(models.Item.scan_data['payload']['row'].astext == query).first() + return ( + db.query(models.Item) + .filter(models.Item.scan_data["payload"]["row"].astext == query) + .first() + ) return db.query(models.Item).filter(models.Item.id == item_id).first() diff --git a/api/main.py b/api/main.py index de59ea9..78f9bed 100644 --- a/api/main.py +++ b/api/main.py @@ -19,13 +19,17 @@ def get_db(): @app.get("/items/", response_model=list[schemas.ItemBase]) -async def read_items(skip: int = 0, limit: int = 100, db: Session = Depends(get_db)) -> List[schemas.ItemBase]: +async def read_items( + skip: int = 0, limit: int = 100, db: Session = Depends(get_db) +) -> List[schemas.ItemBase]: items = crud.get_items(db, skip=skip, limit=limit) return items @app.get("/items/{item_id}", response_model=schemas.ItemBase) -def read_item(item_id: int=0, q: str="", db: Session = Depends(get_db)) -> schemas.ItemBase: +def read_item( + item_id: int = 0, q: str = "", db: Session = Depends(get_db) +) -> schemas.ItemBase: db_item = crud.get_item(db, item_id=item_id, query=q) if db_item is None: raise HTTPException(status_code=404, detail="Item not found") diff --git a/api/schemas.py b/api/schemas.py index 3c89668..6214be3 100644 --- a/api/schemas.py +++ b/api/schemas.py @@ -1,21 +1,23 @@ from pydantic import BaseModel - class Payload(BaseModel): row: str + class Meta(BaseModel): uuid: str ts: str type: str + class ScanData(BaseModel): version: str format: str meta: Meta payload: Payload + class ItemBase(BaseModel): id: int scan_data: ScanData @@ -23,7 +25,6 @@ class ItemBase(BaseModel): class Config: orm_mode = True + class ItemCreate(ItemBase): pass - - diff --git a/pumpkin/correlation.py b/pumpkin/correlation.py index b5e853a..16bf03e 100644 --- a/pumpkin/correlation.py +++ b/pumpkin/correlation.py @@ -57,7 +57,7 @@ async def setup(self): async def main(): jid = "correlation-engine@localhost" - passwd = "securePasswordforCE" # getpass.getpass(f"Password for {jid}:\n") + passwd = "securePasswordforCE" # getpass.getpass(f"Password for {jid}:\n") agent = CorrelationEngine(jid, passwd) await agent.start() @@ -67,7 +67,7 @@ async def main(): print("Wait until user interrupts with ctrl+C") # wait until user interrupts with ctrl+C - while True: # not agent.CollectingBehav.is_killed(): + while True: # not agent.CollectingBehav.is_killed(): try: await asyncio.sleep(1) except KeyboardInterrupt: diff --git a/pumpkin/probe.py b/pumpkin/probe.py index 17d5a6e..47c236d 100644 --- a/pumpkin/probe.py +++ b/pumpkin/probe.py @@ -13,7 +13,7 @@ from spade.message import Message -def exec_cmd(cmd: str="fortune", working_dir: str = "") -> str: +def exec_cmd(cmd: str = "fortune", working_dir: str = "") -> str: """Execute a command in a sub process and wait for the result.""" bash_string = r"""#!/bin/bash set -e @@ -79,13 +79,12 @@ async def setup(self): self.add_behaviour(self.InformBehav) - async def main(probe_jid, passwd): agent = CorrelationEngine(probe_jid, passwd) await agent.start() # wait until user interrupts with ctrl+C - while True: # not agent.CollectingBehav.is_killed(): + while True: # not agent.CollectingBehav.is_killed(): try: await asyncio.sleep(1) except KeyboardInterrupt: @@ -96,7 +95,6 @@ async def main(probe_jid, passwd): await agent.stop() - if __name__ == "__main__": parser = argparse.ArgumentParser(prog="probe-agent") parser.add_argument( @@ -119,4 +117,3 @@ async def main(probe_jid, passwd): passwd = config["passwd"] spade.run(main(probe_jid, passwd)) -