cloudgrep is grep for cloud storage
from core.cloudgrep import CloudGrep
import argparse
import logging

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="CloudGrep searches is grep for cloud storage like S3.")
    parser.add_argument("-b", "--bucket", help="Bucket to search. E.g. my-bucket", required=True)
    parser.add_argument(
        "-q", "--query", help="Text to search for. Will be parsed as a Regex. E.g. example.com", required=True
    )
    parser.add_argument(
        "-p",
        "--prefix",
        help="Optionally filter on the start of the Object name. E.g. logs/",
        required=False,
        default="",
    )
    parser.add_argument(
        "-f", "--filename", help="Optionally filter on Objects that match a keyword. E.g. .log.gz ", required=False
    )
    parser.add_argument(
        "-s",
        "--start_date",
        help="Optionally filter on Objects modified after a Date or Time. E.g. 2022-01-01 ",
        required=False,
    )
    parser.add_argument(
        "-e",
        "--end_date",
        help="Optionally filter on Objects modified before a Date or Time. E.g. 2022-01-01 ",
        required=False,
    )
    parser.add_argument(
        "-fs",
        "--file_size",
        help="Optionally filter on Objects smaller than a file size, in bytes. Defaults to 100 Mb. ",
        required=False,
    )
    parser.add_argument("-d", "--debug", help="Enable Debug logging. ", action="store_true", required=False)
    parser.add_argument(
        "-hf", "--hide_filenames", help="Dont show matching filesnames. ", action="store_true", required=False
    )
    args = vars(parser.parse_args())

    if args["debug"]:
        logging.basicConfig(format="[%(asctime)s]:[%(levelname)s] - %(message)s", level=logging.INFO)

    CloudGrep().search(
        args["bucket"],
        args["query"],
        args["prefix"],
        args["filename"],
        args["start_date"],
        args["end_date"],
        args["file_size"],
        args["hide_filenames"],
    )

cloudgrep.py
cado-security/cloudgrep
import boto3
from datetime import timezone, datetime
from dateutil.parser import parse
import botocore
import concurrent
import tempfile
import re
from typing import Iterator, Optional, List
import logging
import gzip
import zipfile
import os


class CloudGrep:
    def get_all_strings_line(self, file_path: str) -> List[str]:
        """Get all the strings from a file line by line
        We do this instead of f.readlines() as this supports binary files too
        """
        with open(file_path, "rb") as f:
            read_bytes = f.read()
            b = read_bytes.decode("utf-8", "ignore")
            b = b.replace("\n", "\r")
            string_list = b.split("\r")
            return string_list

    def search_file(self, file_name: str, key_name: str, search: str, hide_filenames: bool) -> bool:
        """Regex search of the file line by line"""
        matched = False
        logging.info(f"Searching {file_name} for {search}")
        if key_name.endswith(".gz"):
            with gzip.open(file_name, "rt") as f:
                for line in f:
                    if re.search(search, line):
                        print(f"{key_name}: {line}")
                        matched = True
        elif key_name.endswith(".zip"):
            with tempfile.TemporaryDirectory() as tempdir:
                with zipfile.ZipFile(file_name, "r") as zf:
                    zf.extractall(tempdir)
                    logging.info(f"Extracted {file_name} to {tempdir}")
                    for filename in os.listdir(tempdir):
                        logging.info(f"Searching in zip {filename}")
                        if os.path.isfile(os.path.join(tempdir, filename)):
                            with open(os.path.join(tempdir, filename)) as f:
                                for line in f:
                                    if re.search(search, line):
                                        print(f"{key_name}/{filename}: {line}")
                                        matched = True
        else:
            for line in self.get_all_strings_line(file_name):
                if re.search(search, line):
                    if not hide_filenames:
                        print(f"{key_name}: {line}")
                    else:
                        print(line)
                    matched = True

        return matched

    def download_from_s3_multithread(self, bucket: str, files: List[str], query: str, hide_filenames: bool) -> int:
        """Use ThreadPoolExecutor and boto3 to download every file in the bucket from s3
        Returns number of matched files"""
        client_config = botocore.config.Config(
            max_pool_connections=64,
        )
        matched_count = 0
        s3 = boto3.client("s3", config=client_config)

        # Create a function to download the files
        def download_file(key: str) -> None:
            # Get meta data of file in s3 using boto3
            with tempfile.NamedTemporaryFile() as tmp:
                logging.info(f"Downloading {bucket} {key} to {tmp.name}")
                s3.download_file(bucket, key, tmp.name)
                matched = self.search_file(tmp.name, key, query, hide_filenames)
                if matched:
                    nonlocal matched_count
                    matched_count += 1

        # Use ThreadPoolExecutor to download the files
        with concurrent.futures.ThreadPoolExecutor() as executor:  # type: ignore
            executor.map(download_file, files)
        return matched_count

    def filter_object(
        self,
        obj: dict,
        key_contains: Optional[str],
        from_date: Optional[datetime],
        to_date: Optional[datetime],
        file_size: int,
    ) -> bool:
        last_modified = obj["LastModified"]
        if last_modified and from_date and from_date > last_modified:
            return False  # Object was modified before the from_date
        if last_modified and to_date and last_modified > to_date:
            return False  # Object was modified after the to_date
        if obj["Size"] == 0 and obj["Size"] < file_size:
            return False  # Object is empty or too large
        if key_contains and key_contains not in obj["Key"]:
            return False  # Object does not contain the key_contains string
        return True

    def get_objects(
        self,
        bucket: str,
        prefix: Optional[str],
        key_contains: Optional[str],
        from_date: Optional[datetime],
        end_date: Optional[datetime],
        file_size: int,
    ) -> Iterator[str]:
        """Get all objects in a bucket with a given prefix"""
        s3 = boto3.client("s3")
        paginator = s3.get_paginator("list_objects_v2")
        page_iterator = paginator.paginate(Bucket=bucket, Prefix=prefix)
        for page in page_iterator:
            if "Contents" in page:
                for obj in page["Contents"]:
                    if self.filter_object(obj, key_contains, from_date, end_date, file_size):
                        yield obj["Key"]

    def search(
        self,
        bucket: str,
        query: str,
        prefix: Optional[str] = None,
        key_contains: Optional[str] = None,
        from_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        file_size: int = 100000000,
        hide_filenames: bool = False,
    ) -> None:
        s3_client = boto3.client("s3")
        region = s3_client.get_bucket_location(Bucket=bucket)
        print(
            f"Bucket is in region: {region['LocationConstraint']} : Search from the same region to avoid egress charges."
        )
        parsed_from_date = None
        if from_date:
            parsed_from_date = parse(from_date).astimezone(timezone.utc)  # type: ignore
        parsed_end_date = None
        if end_date:
            parsed_end_date = parse(end_date).astimezone(timezone.utc)  # type: ignore
        matching_keys = list(
            self.get_objects(bucket, prefix, key_contains, parsed_from_date, parsed_end_date, file_size)
        )
        print(f"Searching {len(matching_keys)} files in {bucket} for {query}...")
        self.download_from_s3_multithread(bucket, matching_keys, query, hide_filenames)

core/cloudgrep.py
cado-security/cloudgrep
"""
Basic unit tests for Cloud Grep
python3 -m unittest discover tests
"""
import unittest
import os
import boto3
import timeout_decorator
from moto import mock_s3

from core.cloudgrep import CloudGrep


class CloudGrepTests(unittest.TestCase):
    """Tests for Cloud Grep"""

    def test_true(self) -> None:
        self.assertEqual(1, 1)

    def test_weird_files(self) -> None:
        for filename in os.listdir("./tests/data/"):
            # Just checks we don't crash on any files
            CloudGrep().get_all_strings_line("./tests/data/" + filename)

        self.assertIn("SomeLine", CloudGrep().get_all_strings_line("./tests/data/14_3.log"))

    def test_gzip(self) -> None:
        # Get lines from .gz compressed file
        found = CloudGrep().search_file("./tests/data/000000.gz", "000000.gz", "Running on machine", False)
        self.assertTrue(found)

    def test_zip(self) -> None:
        # Get lines from .zip compressed file
        found = CloudGrep().search_file("./tests/data/000000.zip", "000000.zip", "Running on machine", False)
        self.assertTrue(found)

    @timeout_decorator.timeout(5)  # Normally takes around 3 seconds to run in github actions
    @mock_s3
    def test_e2e(self) -> None:
        # This test uploads a couple of logs to mock s3
        # Then searches them

        _BUCKET = "mybucket"
        _QUERY = "SomeLine"

        conn = boto3.resource("s3", region_name="us-east-1")
        conn.create_bucket(Bucket=_BUCKET)
        s3 = boto3.client("s3", region_name="us-east-1")

        for file_name in ["14_3.log", "35010_7.log", "apache_access.log"]:
            with open(f"./tests/data/{file_name}", "rb") as data:
                s3.upload_fileobj(data, _BUCKET, file_name)

        print("Checking we include every file")
        matching_keys = list(CloudGrep().get_objects(_BUCKET, "", None, None, None, 0))
        assert len(matching_keys) == 3

        print(f"Checking we only get one search hit in: {matching_keys}")
        hits = CloudGrep().download_from_s3_multithread(_BUCKET, matching_keys, _QUERY, False)
        assert hits == 1

        # Upload a log 10 000 times and see how long it takes
        print("Uploading large number of logs")
        for x in range(1000):
            with open("./tests/data/apache_access.log", "rb") as data:
                s3.upload_fileobj(data, _BUCKET, str(x))

        print("Searching")
        CloudGrep().download_from_s3_multithread(_BUCKET, matching_keys, _QUERY, False)
        print("Searched")

tests/tests_unit.py
cado-security/cloudgrep