PrivacyScanData/S3/NewFind/s3recon-master/s3recon/mongodb.py

import re
from enum import Enum
from functools import wraps
from logging import getLogger, basicConfig, INFO

from pymongo import MongoClient
from pymongo.errors import ServerSelectionTimeoutError, PyMongoError

logger = getLogger(__name__)

basicConfig(format="%(message)s", level=INFO)


def with_retry(tries):
    def outer_wrapper(f):
        @wraps(f)
        def inner_wrapper(*args, **kwargs):
            def _retry(t=tries):
                if t <= 0:
                    logger.error("unable to write hit to database")
                    # raise WriteError(f"unable to write to database")
                    return
                try:
                    f(*args, **kwargs)
                except PyMongoError:
                    t -= 1
                    _retry(t)

            return _retry()

        return inner_wrapper

    return outer_wrapper


class StringEnum(Enum):
    def __str__(self):
        return str(self.value["text"])

    def __repr__(self):
        return str(self.value["key"])


class Access(StringEnum):
    PUBLIC = {"key": "+", "text": "public"}
    PRIVATE = {"key": "-", "text": "private"}


class Hit:
    def __init__(self, url: str, access: Access):
        self.url = url
        self.access = access

    def __iter__(self):
        yield from {"url": self.url, "access": str(self.access)}.items()

    def is_valid(self):
        return (
            re.match(r"^https?://.*\.amazonaws.com/.*$", self.url)
            and self.access in Access
        )


class MongoDB:
    def __init__(
        self,
        host: str = "0.0.0.0",
        port: int = 27017,
        db_name: str = "s3recon",
        col_name: str = "hits",
        unique_indicies: tuple = ("url",),
        indicies: tuple = ("access",),
        timeout: int = 10,
    ):
        self.client = MongoClient(host, port, serverSelectionTimeoutMS=timeout)
        self.db_name = db_name
        self.col_name = col_name
        self.index(unique_indicies, unique=True)
        self.index(indicies)

    def __del__(self):
        self.client.close()

    def index(self, indicies=(), **kwargs):
        for i in indicies:
            self.client[self.db_name][self.col_name].ensure_index(i, **kwargs)

    @staticmethod
    def normalize(item):
        if isinstance(item, (list, set)):
            return list(map(dict, item))
        else:
            return dict(item)

    @with_retry(3)
    def insert_many(self, items):
        self.client[self.db_name][self.col_name].insert_many(self.normalize(items))

    @with_retry(3)
    def insert(self, item):
        self.client[self.db_name][self.col_name].insert(self.normalize(item))

    @with_retry(3)
    def update_many(self, filter, items):
        self.client[self.db_name][self.col_name].update_many(
            filter, self.normalize(items), upsert=True
        )

    @with_retry(3)
    def update(self, filter, item):
        self.client[self.db_name][self.col_name].update(
            filter, self.normalize(item), upsert=True
        )

    def is_connected(self):
        try:
            self.client.server_info()
        except ServerSelectionTimeoutError:
            return False
        return True