257 lines
7 KiB
Python
257 lines
7 KiB
Python
"""Store bot state"""
|
|
import os
|
|
from datetime import datetime
|
|
from schema import Schema, And, Use, Optional
|
|
from typing import Protocol
|
|
import files
|
|
import copy
|
|
import logging
|
|
|
|
logger = logging.getLogger("state")
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
# Warning: Changing any of the below needs proper thought to ensure backwards compatibility
|
|
_VERSION = 1
|
|
|
|
|
|
def _migrate_to_v1(d: dict):
|
|
logger.info("Renaming %s to %s", _Key.MATCHEES, _Key.USERS)
|
|
d[_Key.USERS] = d[_Key.MATCHEES]
|
|
del d[_Key.MATCHEES]
|
|
|
|
|
|
# Set of migration functions to apply
|
|
_MIGRATIONS = [
|
|
_migrate_to_v1
|
|
]
|
|
|
|
|
|
class AuthScope(str):
|
|
"""Various auth scopes"""
|
|
OWNER = "owner"
|
|
MATCHER = "matcher"
|
|
|
|
|
|
class _Key(str):
|
|
"""Various keys used in the schema"""
|
|
HISTORY = "history"
|
|
GROUPS = "groups"
|
|
MEMBERS = "members"
|
|
USERS = "users"
|
|
SCOPES = "scopes"
|
|
MATCHES = "matches"
|
|
ACTIVE = "active"
|
|
CHANNELS = "channels"
|
|
REACTIVATE = "reactivate"
|
|
VERSION = "version"
|
|
|
|
# Unused
|
|
MATCHEES = "matchees"
|
|
|
|
|
|
_TIME_FORMAT = "%a %b %d %H:%M:%S %Y"
|
|
|
|
|
|
_SCHEMA = Schema(
|
|
{
|
|
# The current version
|
|
_Key.VERSION: And(Use(int)),
|
|
|
|
Optional(_Key.HISTORY): {
|
|
# A datetime
|
|
Optional(str): {
|
|
_Key.GROUPS: [
|
|
{
|
|
_Key.MEMBERS: [
|
|
# The ID of each matchee in the match
|
|
And(Use(int))
|
|
]
|
|
}
|
|
]
|
|
}
|
|
},
|
|
Optional(_Key.USERS): {
|
|
Optional(str): {
|
|
Optional(_Key.SCOPES): And(Use(list[str])),
|
|
Optional(_Key.MATCHES): {
|
|
# Matchee ID and Datetime pair
|
|
Optional(str): And(Use(str))
|
|
},
|
|
Optional(_Key.CHANNELS): {
|
|
# The channel ID
|
|
Optional(str): {
|
|
# Whether the user is signed up in this channel
|
|
_Key.ACTIVE: And(Use(bool)),
|
|
}
|
|
}
|
|
}
|
|
},
|
|
}
|
|
)
|
|
|
|
# Empty but schema-valid internal dict
|
|
_EMPTY_DICT = {
|
|
_Key.HISTORY: {},
|
|
_Key.USERS: {},
|
|
_Key.VERSION: _VERSION
|
|
}
|
|
assert _SCHEMA.validate(_EMPTY_DICT)
|
|
|
|
|
|
class Member(Protocol):
|
|
@property
|
|
def id(self) -> int:
|
|
pass
|
|
|
|
|
|
def ts_to_datetime(ts: str) -> datetime:
|
|
"""Convert a ts to datetime using the internal format"""
|
|
return datetime.strptime(ts, _TIME_FORMAT)
|
|
|
|
|
|
class State():
|
|
def __init__(self, data: dict = _EMPTY_DICT):
|
|
"""Initialise and validate the state"""
|
|
self.validate(data)
|
|
self._dict = copy.deepcopy(data)
|
|
|
|
@property
|
|
def _history(self) -> dict[str]:
|
|
return self._dict[_Key.HISTORY]
|
|
|
|
@property
|
|
def _users(self) -> dict[str]:
|
|
return self._dict[_Key.USERS]
|
|
|
|
def validate(self, dict: dict = None):
|
|
"""Initialise and validate a state dict"""
|
|
if not dict:
|
|
dict = self._dict
|
|
_SCHEMA.validate(dict)
|
|
|
|
def get_oldest_timestamp(self) -> datetime:
|
|
"""Grab the oldest timestamp in history"""
|
|
times = (ts_to_datetime(dt) for dt in self._history.keys())
|
|
return next(times, None)
|
|
|
|
def get_user_matches(self, id: int) -> list[int]:
|
|
return self._users.get(str(id), {}).get(_Key.MATCHES, {})
|
|
|
|
def log_groups(self, groups: list[list[Member]], ts: datetime = datetime.now()) -> None:
|
|
"""Log the groups"""
|
|
tmp_state = State(self._dict)
|
|
ts = datetime.strftime(ts, _TIME_FORMAT)
|
|
|
|
# Grab or create the hitory item for this set of groups
|
|
history_item = {}
|
|
tmp_state._history[ts] = history_item
|
|
history_item_groups = []
|
|
history_item[_Key.GROUPS] = history_item_groups
|
|
|
|
for group in groups:
|
|
|
|
# Add the group data
|
|
history_item_groups.append({
|
|
_Key.MEMBERS: [m.id for m in group]
|
|
})
|
|
|
|
# Update the matchee data with the matches
|
|
for m in group:
|
|
matchee = tmp_state._users.get(str(m.id), {})
|
|
matchee_matches = matchee.get(_Key.MATCHES, {})
|
|
|
|
for o in (o for o in group if o.id != m.id):
|
|
matchee_matches[str(o.id)] = ts
|
|
|
|
matchee[_Key.MATCHES] = matchee_matches
|
|
tmp_state._users[str(m.id)] = matchee
|
|
|
|
# Validate before storing the result
|
|
tmp_state.validate()
|
|
self._dict = tmp_state._dict
|
|
|
|
def set_user_scope(self, id: str, scope: str, value: bool = True):
|
|
"""Add an auth scope to a user"""
|
|
# Dive in
|
|
user = self._users.get(str(id), {})
|
|
scopes = user.get(_Key.SCOPES, [])
|
|
|
|
# Set the value
|
|
if value and scope not in scopes:
|
|
scopes.append(scope)
|
|
elif not value and scope in scopes:
|
|
scopes.remove(scope)
|
|
|
|
# Roll out
|
|
user[_Key.SCOPES] = scopes
|
|
self._users[id] = user
|
|
|
|
def get_user_has_scope(self, id: str, scope: str) -> bool:
|
|
"""
|
|
Check if a user has an auth scope
|
|
"owner" users have all scopes
|
|
"""
|
|
user = self._users.get(str(id), {})
|
|
scopes = user.get(_Key.SCOPES, [])
|
|
return AuthScope.OWNER in scopes or scope in scopes
|
|
|
|
def set_use_active_in_channel(self, id: str, channel_id: str, active: bool = True):
|
|
"""Set a user as active (or not) on a given channel"""
|
|
# Dive in
|
|
user = self._users.get(str(id), {})
|
|
channels = user.get(_Key.CHANNELS, {})
|
|
channel = channels.get(str(channel_id), {})
|
|
|
|
# Set the value
|
|
channel[_Key.ACTIVE] = active
|
|
|
|
# Unroll
|
|
channels[str(channel_id)] = channel
|
|
user[_Key.CHANNELS] = channels
|
|
self._users[str(id)] = user
|
|
|
|
def get_user_active_in_channel(self, id: str, channel_id: str) -> bool:
|
|
"""Get a users active channels"""
|
|
user = self._users.get(str(id), {})
|
|
channels = user.get(_Key.CHANNELS, {})
|
|
return str(channel_id) in [channel for (channel, props) in channels.items() if props.get(_Key.ACTIVE, False)]
|
|
|
|
@property
|
|
def dict_internal(self) -> dict:
|
|
"""Only to be used to get the internal dict as a copy"""
|
|
return copy.deepcopy(self._dict)
|
|
|
|
|
|
def _migrate(dict: dict):
|
|
"""Migrate a dict through versions"""
|
|
version = dict.get("version", 0)
|
|
for i in range(version, _VERSION):
|
|
logger.info("Migrating from v%s to v%s", version, version+1)
|
|
_MIGRATIONS[i](dict)
|
|
dict[_Key.VERSION] = _VERSION
|
|
|
|
|
|
def load_from_file(file: str) -> State:
|
|
"""
|
|
Load the state from a file
|
|
Apply any required migrations
|
|
"""
|
|
loaded = _EMPTY_DICT
|
|
|
|
# If there's a file load it and try to migrate
|
|
if os.path.isfile(file):
|
|
loaded = files.load(file)
|
|
_migrate(loaded)
|
|
|
|
st = State(loaded)
|
|
|
|
# Save out the migrated (or new) file
|
|
files.save(file, st._dict)
|
|
|
|
return st
|
|
|
|
|
|
def save_to_file(state: State, file: str):
|
|
"""Saves the state out to a file"""
|
|
files.save(file, state.dict_internal)
|