diff --git a/anonstream/control/spec/common.py b/anonstream/control/spec/common.py index bfee339..df4312b 100644 --- a/anonstream/control/spec/common.py +++ b/anonstream/control/spec/common.py @@ -3,9 +3,14 @@ import json -from anonstream.control.spec import Spec, NoParse, Ambiguous, Parsed +from quart import current_app + +from anonstream.control.spec import Spec, NoParse, Ambiguous, BadArgument, Parsed from anonstream.control.spec.utils import get_item, startswith +USERS_BY_TOKEN = current_app.users_by_token +USERS = current_app.users + class Str(Spec): AS_ARG = False @@ -146,3 +151,26 @@ class ArgsJsonStringArray(ArgsJson): f'bad argument at position {index} {obj_json!r}: ' f'could not decode json array of strings' ) + +class ArgsJsonTokenUser(ArgsJsonString): + def transform_obj(self, token): + try: + user = USERS_BY_TOKEN[token] + except KeyError: + raise BadArgument(f'no user with token {token!r}') + return user + +class ArgsJsonHashUser(ArgsString): + def transform_string(self, token_hash): + for user in USERS: + if user['token_hash'] == token_hash: + break + else: + raise BadArgument(f'no user with token_hash {token_hash!r}') + return user + +def ArgsUser(spec): + return Str({ + 'token': ArgsJsonTokenUser(spec), + 'hash': ArgsJsonHashUser(spec), + }) diff --git a/anonstream/control/spec/methods/user.py b/anonstream/control/spec/methods/user.py index 88af4bb..5cd91fd 100644 --- a/anonstream/control/spec/methods/user.py +++ b/anonstream/control/spec/methods/user.py @@ -6,40 +6,15 @@ import json from quart import current_app from anonstream.control.exceptions import CommandFailed -from anonstream.control.spec import BadArgument -from anonstream.control.spec.common import Str, End, ArgsInt, ArgsString, ArgsJson, ArgsJsonBoolean, ArgsJsonString -from anonstream.control.spec.utils import get_item, json_dumps_contiguous +from anonstream.control.spec.common import Str, End, ArgsInt, ArgsString, ArgsJson, ArgsJsonBoolean, ArgsJsonString, ArgsUser +from anonstream.control.spec.utils import json_dumps_contiguous from anonstream.utils.user import USER_WEBSOCKET_ATTRS from anonstream.routes.wrappers import generate_and_add_user from anonstream.wrappers import get_timestamp USERS_BY_TOKEN = current_app.users_by_token -USERS = current_app.users USERS_UPDATE_BUFFER = current_app.users_update_buffer -class ArgsJsonTokenUser(ArgsJsonString): - def transform_obj(self, token): - try: - user = USERS_BY_TOKEN[token] - except KeyError: - raise BadArgument(f'no user with token {token!r}') - return user - -class ArgsJsonHashUser(ArgsString): - def transform_string(self, token_hash): - for user in USERS: - if user['token_hash'] == token_hash: - break - else: - raise BadArgument(f'no user with token_hash {token_hash!r}') - return user - -def ArgsUser(spec): - return Str({ - 'token': ArgsJsonTokenUser(spec), - 'hash': ArgsJsonHashUser(spec), - }) - async def cmd_user_help(): normal = ['user', 'help'] response = (