diff --git a/.gitignore b/.gitignore index 1b019739..b7e3188b 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ config.yaml registration.yaml *.db *.session +*.json diff --git a/mautrix_appservice/appservice.py b/mautrix_appservice/appservice.py index 8a25e5e9..0b66d57b 100644 --- a/mautrix_appservice/appservice.py +++ b/mautrix_appservice/appservice.py @@ -34,7 +34,8 @@ class AppService: self.as_token = as_token self.hs_token = hs_token self.bot_mxid = f"@{bot_localpart}:{domain}" - self.state_store = StateStore() + self.state_store = StateStore(autosave_file="mx-state.json") + self.state_store.load("mx-state.json") self.transactions = [] diff --git a/mautrix_appservice/intent_api.py b/mautrix_appservice/intent_api.py index 9de0bbeb..aa8b322c 100644 --- a/mautrix_appservice/intent_api.py +++ b/mautrix_appservice/intent_api.py @@ -152,7 +152,6 @@ class IntentAPI: self.localpart = results.group(1) self.state_store = state_store - self.registered = False def user(self, user): if not self.bot: @@ -336,15 +335,16 @@ class IntentAPI: raise IntentError(f"Failed to join room {room_id} as {self.mxid}", e2) def _ensure_registered(self): - if self.registered: + if self.state_store.is_registered(self.mxid): return try: self.client.register({"username": self.localpart}) except MatrixRequestError as e: if matrix_error_code(e) != "M_USER_IN_USE": self.log.exception(f"Failed to register {self.mxid}!") + return # raise IntentError(f"Failed to register {self.mxid}", e) - self.registered = True + self.state_store.registered(self.mxid) def _ensure_has_power_level_for(self, room_id, event_type): if not self.state_store.has_power_level_data(room_id): diff --git a/mautrix_appservice/state_store.py b/mautrix_appservice/state_store.py index bdfa211a..decc35d5 100644 --- a/mautrix_appservice/state_store.py +++ b/mautrix_appservice/state_store.py @@ -14,11 +14,61 @@ # # You should have received a copy of the GNU General Public License # along with this program. If not, see . +import json + class StateStore: - def __init__(self): + def __init__(self, autosave_file=None): + self.registrations = set() self.memberships = {} self.power_levels = {} + self.autosave_file = autosave_file + + def save(self, file): + if isinstance(file, str): + output = open(file, "w") + else: + output = file + + json.dump({ + "registrations": list(self.registrations), + "memberships": self.memberships, + "power_levels": self.power_levels, + }, output) + + if isinstance(file, str): + output.close() + + def load(self, file): + if isinstance(file, str): + try: + input = open(file, "r") + except FileNotFoundError: + return + else: + input = file + + data = json.load(input) + if "registrations" in data: + self.registrations = set(data["registrations"]) + if "memberships" in data: + self.memberships = data["memberships"] + if "power_levels" in data: + self.power_levels = data["power_levels"] + + if isinstance(file, str): + input.close() + + def _autosave(self): + if self.autosave_file: + self.save(self.autosave_file) + + def is_registered(self, user): + return user in self.registrations + + def registered(self, user): + self.registrations.add(user) + self._autosave() def _get_membership(self, room, user): return self.memberships.get(room, {}).get(user, "left") @@ -30,6 +80,7 @@ class StateStore: if room not in self.memberships: self.memberships[room] = {} self.memberships[room][user] = membership + self._autosave() def joined(self, room, user): return self._set_membership(room, user, "join") @@ -56,6 +107,7 @@ class StateStore: "events": {}, } self.power_levels[room]["users"][user] = level + self._autosave() def set_power_levels(self, room, content): if "events" not in content: @@ -63,3 +115,4 @@ class StateStore: if "users" not in content: content["users"] = {} self.power_levels[room] = content + self._autosave()