Пример #1
0
def setup_db_before_test():
    """Sets up database automatically before each test"""
    _db = SqliteDatabase(":memory:")
    with _db.bind_ctx(MODELS):
        _db.create_tables(MODELS)
        yield _db
        _db.drop_tables(MODELS)
 def test_db_closure(*args, **kwargs):
     test_db = SqliteDatabase(":memory:")
     with test_db.bind_ctx(dbs):
         test_db.create_tables(dbs)
         try:
             func(*args, **kwargs)
         finally:
             test_db.drop_tables(dbs)
             test_db.close()
Пример #3
0
def clean_database():
    old_db = database.db
    try:
        test_db = SqliteDatabase(':memory:')
        database.db = test_db
        with test_db.bind_ctx(database.all_classes):
            test_db.connect(reuse_if_open=True)
            test_db.create_tables(database.all_classes)
            yield test_db
    finally:
        database.db = old_db
def test_initialize() -> None:
    test_db = SqliteDatabase(":memory:")
    with test_db.bind_ctx(MODELS):
        initialize(test_db)
        assert test_db.get_tables() == [t._meta.table_name for t in MODELS]
Пример #5
0
class LegacyMatrixStore(object):
    """Storage class for matrix state."""

    models = [
        LegacyAccounts,
        LegacyOlmSessions,
        LegacyMegolmInboundSessions,
        LegacyForwardedChains,
        LegacyDeviceKeys,
        LegacyEncryptedRooms,
        LegacyOutgoingKeyRequests,
    ]

    user_id = attr.ib(type=str)
    device_id = attr.ib(type=str)
    store_path = attr.ib(type=str)
    pickle_key = attr.ib(type=str, default="")
    database_name = attr.ib(type=str, default="")
    database_path = attr.ib(type=str, init=False)
    database = attr.ib(type=SqliteDatabase, init=False)

    def __attrs_post_init__(self):
        self.database_name = self.database_name or "{}_{}.db".format(
            self.user_id, self.device_id)
        self.database_path = os.path.join(self.store_path, self.database_name)
        self.database = SqliteDatabase(self.database_path,
                                       pragmas={
                                           "foreign_keys": 1,
                                           "secure_delete": 1,
                                       })
        with self.database.bind_ctx(self.models):
            self.database.connect(reuse_if_open=True)
            self.database.create_tables(self.models)

    @use_database
    def close(self):
        self.database.close()

    @use_database
    def _get_account(self):
        try:
            return LegacyAccounts.get(
                LegacyAccounts.user_id == self.user_id,
                LegacyAccounts.device_id == self.device_id)
        except DoesNotExist:
            return None

    def load_account(self):
        # type: () -> Optional[OlmAccount]
        """Load the Olm account from the database.

        Returns:
            ``OlmAccount`` object, or ``None`` if it wasn't found for the
                current device_id.

        """
        account = self._get_account()

        if not account:
            return None

        return OlmAccount.from_pickle(account.account, self.pickle_key,
                                      account.shared)

    @use_database
    def save_account(self, account):
        """Save the provided Olm account to the database.

        Args:
            account (OlmAccount): The olm account that will be pickled and
                saved in the database.
        """
        LegacyAccounts.insert(
            user_id=self.user_id,
            device_id=self.device_id,
            shared=account.shared,
            account=account.pickle(
                self.pickle_key)).on_conflict_ignore().execute()

        LegacyAccounts.update({
            LegacyAccounts.account:
            account.pickle(self.pickle_key),
            LegacyAccounts.shared:
            account.shared
        }).where((LegacyAccounts.user_id == self.user_id)
                 & (LegacyAccounts.device_id == self.device_id)).execute()

    @use_database
    def load_sessions(self):
        # type: () -> SessionStore
        """Load all Olm sessions from the database.

        Returns:
            ``SessionStore`` object, containing all the loaded sessions.

        """
        session_store = SessionStore()

        sessions = LegacyOlmSessions.select().join(LegacyAccounts).where(
            LegacyAccounts.device_id == self.device_id)

        for s in sessions:
            session = Session.from_pickle(s.session, s.creation_time,
                                          self.pickle_key)
            session_store.add(s.curve_key, session)

        return session_store

    @use_database
    def save_session(self, curve_key, session):
        """Save the provided Olm session to the database.

        Args:
            curve_key (str): The curve key that owns the Olm session.
            session (Session): The Olm session that will be pickled and
                saved in the database.
        """
        LegacyOlmSessions.replace(
            device=self.device_id,
            curve_key=curve_key,
            session=session.pickle(self.pickle_key),
            session_id=session.id,
            creation_time=session.creation_time).execute()

    @use_database
    def load_inbound_group_sessions(self):
        # type: () -> GroupSessionStore
        """Load all Olm sessions from the database.

        Returns:
            ``GroupSessionStore`` object, containing all the loaded sessions.

        """
        store = GroupSessionStore()

        sessions = LegacyMegolmInboundSessions.select().join(
            LegacyAccounts).where(LegacyAccounts.device_id == self.device_id)

        for s in sessions:
            session = InboundGroupSession.from_pickle(
                s.session, s.ed_key, s.curve_key, s.room_id, self.pickle_key,
                [chain.curve_key for chain in s.forwarded_chains])
            store.add(session)

        return store

    @use_database
    def save_inbound_group_session(self, session):
        """Save the provided Megolm inbound group session to the database.

        Args:
            session (InboundGroupSession): The session to save.
        """
        LegacyMegolmInboundSessions.insert(
            curve_key=session.sender_key,
            device=self.device_id,
            ed_key=session.ed25519,
            room_id=session.room_id,
            session=session.pickle(self.pickle_key),
            session_id=session.id).on_conflict_ignore().execute()

        LegacyMegolmInboundSessions.update({
            LegacyMegolmInboundSessions.session:
            session.pickle(self.pickle_key)
        }).where(
            LegacyMegolmInboundSessions.session_id == session.id).execute()

        # TODO, use replace many here
        for chain in session.forwarding_chain:
            LegacyForwardedChains.replace(curve_key=chain,
                                          session=session.id).execute()

    @use_database
    def load_device_keys(self):
        # type: () -> DeviceStore
        """Load all the device keys from the database.

        Returns DeviceStore containing the OlmDevices with the device keys.
        """
        store = DeviceStore()
        device_keys = LegacyDeviceKeys.select().join(LegacyAccounts).where(
            LegacyAccounts.device_id == self.device_id)

        for d in device_keys:
            store.add(
                OlmDevice(
                    d.user_id,
                    d.user_device_id,
                    {
                        "ed25519": d.ed_key,
                        "curve25519": d.curve_key
                    },
                    display_name="",
                    deleted=d.deleted,
                ))

        return store

    @use_database_atomic
    def save_device_keys(self, device_keys):
        """Save the provided device keys to the database.

        Args:
            device_keys (Dict[str, Dict[str, OlmDevice]]): A dictionary
                containing a mapping from a user id to a dictionary containing
                a mapping of a device id to a OlmDevice.
        """
        rows = []

        for user_id, devices_dict in device_keys.items():
            for device_id, device in devices_dict.items():
                rows.append({
                    "curve_key": device.curve25519,
                    "deleted": device.deleted,
                    "device": self.device_id,
                    "ed_key": device.ed25519,
                    "user_device_id": device_id,
                    "user_id": user_id,
                })

        if not rows:
            return

        for idx in range(0, len(rows), 100):
            data = rows[idx:idx + 100]
            LegacyDeviceKeys.replace_many(data).execute()

    @use_database
    def load_encrypted_rooms(self):
        """Load the set of encrypted rooms for this account.

        Returns:
            ``Set`` containing room ids of encrypted rooms.

        """
        account = self._get_account()

        if not account:
            return set()

        return {room.room_id for room in account.encrypted_rooms}

    @use_database
    def load_outgoing_key_requests(self):
        """Load the set of outgoing key requests for this account.

        Returns:
            ``Set`` containing request ids of key requests.

        """
        account = self._get_account()

        if not account:
            return dict()

        return {
            request.request_id: OutgoingKeyRequest.from_response(request)
            for request in account.key_requests
        }

    @use_database
    def add_outgoing_key_request(self, key_request):
        # type: (OutgoingKeyRequest) -> None
        """Add and store a outgoing key request to the store."""
        account = self._get_account()
        assert account

        LegacyOutgoingKeyRequests.insert(
            request_id=key_request.request_id,
            session_id=key_request.session_id,
            room_id=key_request.room_id,
            algorithm=key_request.algorithm,
            device=account.device_id).on_conflict_ignore().execute()

    @use_database_atomic
    def save_encrypted_rooms(self, rooms):
        """Save the set of room ids for this account."""
        account = self._get_account()

        assert account

        data = [(room_id, account) for room_id in rooms]

        for idx in range(0, len(data), 400):
            rows = data[idx:idx + 400]
            LegacyEncryptedRooms.insert_many(rows,
                                             fields=[
                                                 LegacyEncryptedRooms.room_id,
                                                 LegacyEncryptedRooms.account
                                             ]).on_conflict_ignore().execute()

    def blacklist_device(self, device):
        # type: (OlmDevice) -> bool
        raise NotImplementedError

    def unblacklist_device(self, device):
        # type: (OlmDevice) -> bool
        raise NotImplementedError

    def verify_device(self, device):
        # type: (OlmDevice) -> bool
        raise NotImplementedError

    def is_device_verified(self, device):
        # type: (OlmDevice) -> bool
        raise NotImplementedError

    def is_device_blacklisted(self, device):
        # type: (OlmDevice) -> bool
        raise NotImplementedError

    def unverify_device(self, device):
        # type: (OlmDevice) -> bool
        raise NotImplementedError

    def ignore_device(self, device):
        # type: (OlmDevice) -> bool
        """Mark a device as ignored.

        Args:
            device (OlmDevice): The device that will be ignored.

        Returns True if the device has been ignored, False if it already has
        been ignored before.
        """
        raise NotImplementedError

    def unignore_device(self, device):
        # type: (OlmDevice) -> bool
        """Unmark a device as ignored.

        Args:
            device (OlmDevice): The device that will be unignored.

        Returns True if the device has been unignored, False if it wasn't
        ignored in the first place.
        """
        raise NotImplementedError

    def is_device_ignored(self, device):
        # type: (OlmDevice) -> bool
        """Check if a device is ignored.

        Args:
            device (OlmDevice): The device that will be checked if it is
                ignored.

        Returns True if the device is ignored, False otherwise.
        """
        raise NotImplementedError

    def ignore_devices(self, devices):
        # type: (List[OlmDevice]) -> None
        """Mark multiple devices as ignored.

        Args:
            devices (List[OlmDevice]): The devices that will be makred as
                ignored.
        """
        raise NotImplementedError
Пример #6
0
def test_init(tmpdir):
    test_db = SqliteDatabase(':memory:')
    with test_db.bind_ctx(database.all_classes):
        test_db.connect(reuse_if_open=True)
        test_db.create_tables(database.all_classes)