diff --git a/.github/workflows/push_hub.yml b/.github/workflows/push_hub.yml index b56489d2..8f156aa1 100644 --- a/.github/workflows/push_hub.yml +++ b/.github/workflows/push_hub.yml @@ -41,3 +41,20 @@ jobs: repository: mqueryci/mquery-daemon tags: ${{ github.sha }} push: ${{ github.event_name == 'push' }} + build_indexer: + name: Build image + runs-on: ubuntu-latest + env: + DOCKER_BUILDKIT: 1 + steps: + - name: Check out repository + uses: actions/checkout@v2 + - name: Build and push the image + uses: docker/build-push-action@v1.1.0 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + dockerfile: ./deploy/docker/indexer.Dockerfile + repository: mqueryci/mquery-indexer + tags: ${{ github.sha }} + push: ${{ github.event_name == 'push' }} diff --git a/deploy/docker/dev.daemon.Dockerfile b/deploy/docker/dev.daemon.Dockerfile index d6b7027e..a6bc4023 100644 --- a/deploy/docker/dev.daemon.Dockerfile +++ b/deploy/docker/dev.daemon.Dockerfile @@ -9,4 +9,4 @@ RUN pip install watchdog COPY requirements.txt src/plugins/requirements-*.txt /tmp/ RUN ls /tmp/requirements*.txt | xargs -i,, pip --no-cache-dir install -r ,, -CMD pip install -e /usr/src/app && watchmedo auto-restart --pattern=*.py --recursive -- mquery-daemon --with-indexer +CMD pip install -e /usr/src/app && watchmedo auto-restart --pattern=*.py --recursive -- mquery-daemon diff --git a/deploy/docker/dev.indexer.Dockerfile b/deploy/docker/dev.indexer.Dockerfile new file mode 100644 index 00000000..db1a4301 --- /dev/null +++ b/deploy/docker/dev.indexer.Dockerfile @@ -0,0 +1,12 @@ +FROM python:3.10 + +WORKDIR /usr/src/app/src + +RUN apt update; apt install -y cmake + +# mquery and plugin requirements +RUN pip install watchdog +COPY requirements.txt src/plugins/requirements-*.txt /tmp/ +RUN ls /tmp/requirements*.txt | xargs -i,, pip --no-cache-dir install -r ,, + +CMD pip install -e /usr/src/app && watchmedo auto-restart --pattern=*.py --recursive -- mquery-indexer diff --git a/deploy/docker/indexer.Dockerfile b/deploy/docker/indexer.Dockerfile new file mode 100644 index 00000000..798512b4 --- /dev/null +++ b/deploy/docker/indexer.Dockerfile @@ -0,0 +1,13 @@ +FROM python:3.10 + +RUN apt update; apt install -y cmake + +# mquery and plugin requirements +COPY requirements.txt src/plugins/requirements-*.txt /tmp/ +RUN ls /tmp/requirements*.txt | xargs -i,, pip --no-cache-dir install -r ,, + +COPY requirements.txt setup.py MANIFEST.in /app/ +COPY src /app/src/ +RUN pip install /app + +ENTRYPOINT ["mquery-indexer"] diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index f29eeabb..ddfe2284 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -47,6 +47,31 @@ services: volumes: - "${SAMPLES_DIR}:/mnt/samples" - .:/usr/src/app + depends_on: + dev-web: + condition: service_healthy + redis: + condition: service_started + ursadb: + condition: service_started + postgres: + condition: service_healthy + environment: + - "REDIS_HOST=redis" + - "MQUERY_BACKEND=tcp://ursadb:9281" + - "MQUERY_PLUGINS=${MQUERY_PLUGINS}" + - "DATABASE_URL=postgresql://postgres:password@postgres:5432/mquery" + dev-indexer: + build: + context: . + dockerfile: deploy/docker/dev.indexer.Dockerfile + links: + - redis + - ursadb + - postgres + volumes: + - "${SAMPLES_DIR}:/mnt/samples" + - .:/usr/src/app - "s3:/mnt/s3" depends_on: dev-web: diff --git a/setup.py b/setup.py index 6668a4ab..c35b8197 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,7 @@ install_requires=open("requirements.txt").read().splitlines(), scripts=[ "src/scripts/mquery-daemon", + "src/scripts/mquery-indexer", ], classifiers=[ "Programming Language :: Python", diff --git a/src/app.py b/src/app.py index b8c43c0f..f66fe071 100644 --- a/src/app.py +++ b/src/app.py @@ -310,7 +310,7 @@ def backend_status_datasets() -> BackendStatusDatasetsSchema: for agent_spec in db.get_active_agents().values(): try: ursa = UrsaDb(agent_spec.ursadb_url) - datasets.update(ursa.topology()["result"]["datasets"]) + datasets.update(ursa.datasets()) except Again: pass @@ -616,17 +616,6 @@ def delete_queued_by_id(ursadb_id: str) -> StatusSchema: return StatusSchema(status="ok") -@app.post( - "/api/queue/{ursadb_id}/index", - response_model=StatusSchema, - tags=["queue"], - dependencies=[Depends(can_manage_queues)], -) -def index_queue(ursadb_id: str) -> StatusSchema: - db.create_indexing_task(ursadb_id) - return StatusSchema(status="ok") - - # Permissionless routes. # 1. Static routes are always publicly accessible without authorisation. # 2. /api/server is a special route always accessible for everyone. diff --git a/src/daemon.py b/src/daemon.py index 09f10271..5de48b1e 100644 --- a/src/daemon.py +++ b/src/daemon.py @@ -21,15 +21,6 @@ def start_worker(args: argparse.Namespace, process_index: int) -> None: w.work() -def start_indexer(args: argparse.Namespace) -> None: - setup_logging() - logging.info("Indexer [%s] running...", args.group_id) - - with Connection(Redis(app_config.redis.host, app_config.redis.port)): - w = Worker([args.group_id + ":indexer"]) - w.work() - - def main() -> None: """Spawns a new agent process. Use argv if you want to use a different group_id (it's `default` by default). @@ -48,11 +39,6 @@ def main() -> None: help="Specifies the number of concurrent workers to use for yara matching.", default=1, ) - parser.add_argument( - "--with-indexer", - action="store_true", - help="If specified, indexing worker will be created (not affected by scale).", - ) args = parser.parse_args() @@ -63,8 +49,6 @@ def main() -> None: children = [ Process(target=start_worker, args=(args, i)) for i in range(args.scale) ] - if args.with_indexer: - children.append(Process(target=start_indexer, args=(args,))) for child in children: child.start() diff --git a/src/db.py b/src/db.py index c7420552..d6cf9a6a 100644 --- a/src/db.py +++ b/src/db.py @@ -347,14 +347,6 @@ def create_search_task( self.__schedule(agent, tasks.start_search, job) return job - def create_indexing_task(self, ursadb_id: str) -> None: - """Asks indexer for `ursadb_id` to index pending files.""" - from . import tasks - - Queue(f"{ursadb_id}:indexer", connection=self.redis).enqueue( - tasks.index_pending_files - ) - def get_job_matches( self, job_id: JobId, offset: int = 0, limit: Optional[int] = None ) -> MatchesSchema: @@ -531,9 +523,9 @@ def delete_queued_files(self, ursadb_id: str) -> None: session.query(QueuedFile).filter_by(ursadb_id=ursadb_id).delete() session.commit() - def get_indexing_batch(self, ursadb_id: str) -> List[QueuedFile]: - BATCH_SIZE = 100 - + def get_pending_files( + self, ursadb_id: str, limit: int + ) -> List[QueuedFile]: # We can only batch files with the same tags and index types, so first find # the biggest group of files waiting to be indexed: group_subquery = ( @@ -548,7 +540,7 @@ def get_indexing_batch(self, ursadb_id: str) -> List[QueuedFile]: .subquery() ) - # Then return up to BATCH_SIZE files from that group. + # Then return files from that group. with self.session() as session: group = session.execute(select(group_subquery)).first() if not group: @@ -561,12 +553,12 @@ def get_indexing_batch(self, ursadb_id: str) -> List[QueuedFile]: QueuedFile.index_types == group.index_types, QueuedFile.tags == group.tags, ) - .limit(BATCH_SIZE) + .limit(limit) ).all() return files - def complete_indexing_batch(self, files: list[QueuedFile]) -> None: + def remove_from_pending(self, files: list[QueuedFile]) -> None: ids = [f.id for f in files] with self.session() as session: session.execute( diff --git a/src/e2etests/test_api.py b/src/e2etests/test_api.py index 6f49cf99..409af906 100644 --- a/src/e2etests/test_api.py +++ b/src/e2etests/test_api.py @@ -136,7 +136,7 @@ def test_query_with_taints(add_files_to_index): db = UrsaDb("tcp://ursadb:9281") random_taint = os.urandom(8).hex() - for dataset_id in db.topology()["result"]["datasets"].keys(): + for dataset_id in db.datasets().keys(): out = db.execute_command( f'dataset "{dataset_id}" taint "{random_taint}";' ) diff --git a/src/indexer.py b/src/indexer.py new file mode 100644 index 00000000..bb6402df --- /dev/null +++ b/src/indexer.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python +import argparse +import logging +from multiprocessing import Pool +from time import sleep + +from .db import Database +from .plugins import PluginManager +from .util import setup_logging +from .models.queuedfile import QueuedFile +from .lib.ursadb import UrsaDb +from .config import app_config + + +COMPACT_THRESHOLD = 0 +"""Global variable dependeng on number of workers""" + + +def index_batch( + job: tuple[list[QueuedFile], list[str], list[str]], +) -> list[QueuedFile]: + batch, index_types, tags = job + assert ( + batch + ), "This function shouldn't be called without files, likely a bug." + + paths = [f.path for f in batch] + + db = Database(app_config.redis.host, app_config.redis.port) + ursa = UrsaDb(app_config.mquery.backend) + plugins = PluginManager(app_config.mquery.plugins, db) + + current_datasets = len(ursa.datasets()) + if current_datasets > COMPACT_THRESHOLD: + ursa.execute_command("compact smart;") + + ursadb_batch = [] + for file_path in paths: + final_path = plugins.filter(file_path) + if final_path is None: + logging.debug("Filtering out file %s", file_path) + continue + ursadb_batch.append(final_path) + + logging.debug("Batch preprocessed, asking ursadb to index it.") + + ursa.index(ursadb_batch, index_types=index_types, tags=tags) + logging.debug("Ursadb indexing completed") + + plugins.cleanup() + logging.debug("Cleanup completed") + + return batch + + +def indexer_main(group_id: str, scale: int) -> None: + """Do the indexing in an infinite loop.""" + logging.info("Starting indexer for group %s", group_id) + logging.info("Workers: %s", scale) + + global COMPACT_THRESHOLD + COMPACT_THRESHOLD = scale * 20 + 40 + logging.info("Compact threshold: %s", COMPACT_THRESHOLD) + + db = Database(app_config.redis.host, app_config.redis.port) + + # How many files should one worker index at once + BATCH_SIZE = 1000 + + # Don't get all the files from the database at once, to avoid huge queries + LIMIT = scale * 100 * BATCH_SIZE + + while True: + pending = db.get_pending_files(group_id, LIMIT) + if not pending: + logging.debug("No pending files left.") + sleep(15) # Wait a bit to collect some files. + continue + + index_types, tags = pending[0].index_types, pending[0].tags + logging.debug("We are indexing type=%s, tags=%s", index_types, tags) + + batches = [] + next_batch = [] + for f in pending: + next_batch.append(f) + if len(next_batch) > BATCH_SIZE: + batches.append(next_batch) + next_batch = [] + + if next_batch: + batches.append(next_batch) + + types_str = ",".join(type for type in index_types) + tags_str = ",".join(tag for tag in tags) + logging.info( + "[0/%s] Collected batches (%s files), with [%s], with tags [%s]", + len(batches), + len(pending), + types_str, + tags_str, + ) + + jobs = [(batch, index_types, tags) for batch in batches] + + done = 0 + pool = Pool(processes=scale) + for completed_batch in pool.imap_unordered( + index_batch, jobs, chunksize=1 + ): + done += 1 + db.remove_from_pending(completed_batch) + logging.info("[%s/%s] Batch completed.", done, len(batches)) + + logging.info("[%s/%s] Indexing completed", done, len(batches)) + + +def main() -> None: + """Spawns a new indexer process. Indexer will work in the background all the time.""" + + parser = argparse.ArgumentParser( + description="Start mquery indexer worker." + ) + parser.add_argument( + "group_id", + help="Name of the agent group to join to", + nargs="?", + default="default", + ) + parser.add_argument( + "--scale", + type=int, + help="Specifies the number of concurrent workers to use for yara matching.", + default=1, + ) + args = parser.parse_args() + + setup_logging(logging.DEBUG) + indexer_main(args.group_id, args.scale) + + +if __name__ == "__main__": + main() diff --git a/src/lib/ursadb.py b/src/lib/ursadb.py index 01bb10e6..22027a8f 100644 --- a/src/lib/ursadb.py +++ b/src/lib/ursadb.py @@ -45,7 +45,7 @@ def __init__(self, backend: str) -> None: self.backend = backend def __execute(self, command: str, recv_timeout: int = 2000) -> Json: - logging.info("Ursadb command: %s", command) + logging.debug("Ursadb command: %s", command) context = zmq.Context() try: socket = context.socket(zmq.REQ) @@ -105,7 +105,15 @@ def status(self) -> Json: return self.__execute("status;") def topology(self) -> Json: - return self.__execute("topology;") + result = self.__execute("topology;") + + if "error" in result: + raise RuntimeError(result["error"]) + + return result["result"] + + def datasets(self) -> Json: + return self.topology()["datasets"] def execute_command(self, command: str) -> Json: return self.__execute(command, -1) diff --git a/src/plugins/__init__.py b/src/plugins/__init__.py index f00754ed..08bbd88c 100644 --- a/src/plugins/__init__.py +++ b/src/plugins/__init__.py @@ -48,7 +48,7 @@ def __init__(self, spec: str, db: Database) -> None: plugin_config = db.get_plugin_config(plugin_name) try: active_plugins.append(plugin_class(db, plugin_config)) - logging.info("Loaded plugin %s", plugin_name) + logging.debug("Loaded plugin %s", plugin_name) except Exception: logging.exception("Failed to load %s plugin", plugin_name) self.active_plugins = active_plugins diff --git a/src/scripts/mquery-indexer b/src/scripts/mquery-indexer new file mode 100644 index 00000000..d5acde3d --- /dev/null +++ b/src/scripts/mquery-indexer @@ -0,0 +1,6 @@ +#!/usr/bin/env python + +import mquery.indexer + + +mquery.indexer.main() diff --git a/src/tasks.py b/src/tasks.py index 97ebc741..eebdacb0 100644 --- a/src/tasks.py +++ b/src/tasks.py @@ -12,7 +12,6 @@ from .plugins import PluginManager from .models.job import Job, JobStatus from .models.match import Match -from .models.queuedfile import QueuedFile from .lib.yaraparse import parse_yara, combine_rules from .lib.ursadb import Json, UrsaDb from .metadata import Metadata @@ -62,12 +61,7 @@ def register(self) -> None: def get_datasets(self) -> List[str]: """Returns a list of dataset IDs, or throws an exception on error.""" - result = self.ursa.topology() - - if "error" in result: - raise RuntimeError(result["error"]) - - return list(result["result"]["datasets"].keys()) + return list(self.ursa.datasets().keys()) def update_metadata( self, @@ -339,56 +333,3 @@ def read_bytes_with_context( matching = data[offset : offset + length] after = data[offset + length : offset + length + context] return before, matching, after - - -def __index_single_batch( - plugins: PluginManager, files_to_index: list[QueuedFile] -) -> None: - ursa = UrsaDb(app_config.mquery.backend) - ursadb_batch = [] - - if not files_to_index: - logging.info("Nothing to index") - return - - paths = [f.path for f in files_to_index] - representative = files_to_index[0] - index_types = representative.index_types - tags = representative.tags - - logging.debug("Raw files to index %s", paths) - logging.info("Preprocessing batch of size: %s", len(paths)) - for file_path in paths: - final_path = plugins.filter(file_path) - if final_path is None: - logging.debug("Filtering out file %s", file_path) - continue - ursadb_batch.append(final_path) - - logging.debug("Final paths %s", ursadb_batch) - logging.info("Batch preprocessed, asking ursadb to index it.") - - ursa.index(ursadb_batch, index_types=index_types, tags=tags) - plugins.cleanup() - - -def index_pending_files() -> None: - db = Database(app_config.redis.host, app_config.redis.port) - plugins = PluginManager(app_config.mquery.plugins, db) - - job = get_current_job() - assert ( - job is not None - ), "index_pending_files called outside of a job context" - - group_id = job.origin.replace(":indexer", "") - logging.info("Starting indexing job for group %s", group_id) - - while True: - batch = db.get_indexing_batch(group_id) - if not batch: - logging.info("Indexing job done for group %s", group_id) - break - - __index_single_batch(plugins, batch) - db.complete_indexing_batch(batch) diff --git a/src/util.py b/src/util.py index 4f5a154f..cd9bff99 100644 --- a/src/util.py +++ b/src/util.py @@ -7,10 +7,8 @@ LOG_DATEFMT = "%Y-%m-%d %H:%M:%S" -def setup_logging() -> None: - logging.basicConfig( - level=logging.INFO, format=LOG_FORMAT, datefmt=LOG_DATEFMT - ) +def setup_logging(level=logging.INFO) -> None: + logging.basicConfig(level=level, format=LOG_FORMAT, datefmt=LOG_DATEFMT) def mquery_version() -> str: diff --git a/src/utils/index.py b/src/utils/index.py index 3a211f95..5b087944 100644 --- a/src/utils/index.py +++ b/src/utils/index.py @@ -59,7 +59,7 @@ def find_new_files( def index_files( - proc_params: Tuple[str, List[str], List[str], Path, int] + proc_params: Tuple[str, List[str], List[str], Path, int], ) -> str: ursa_url, types, tags, batch, compact_threshold = proc_params ursa = UrsaDb(ursa_url)