diff --git a/manage/cmd/migration/command.py b/manage/cmd/migration/command.py index a1e87c6..4f80b64 100644 --- a/manage/cmd/migration/command.py +++ b/manage/cmd/migration/command.py @@ -1,4 +1,5 @@ import inspect +import os from pathlib import Path from dataclasses import dataclass from collections import namedtuple @@ -22,14 +23,13 @@ class MigrationContext: @property def latest(self): """Return the latest migration ID.""" - return max(self.scripts.keys()) - + return 0 if len(self.scripts) == 0 else max(self.scripts.keys()) def make_migration_ctx() -> MigrationContext: """Create the MigrationContext instance.""" # taken from https://stackoverflow.com/a/6628348 script_path = inspect.stack()[0][1] - script_folder = '/'.join(script_path.split('/')[:-1]) + script_folder = os.sep.join(script_path.split(os.sep)[:-1]) script_folder = Path(script_folder) migration_folder = script_folder / 'scripts' @@ -40,7 +40,7 @@ def make_migration_ctx() -> MigrationContext: mig_path_str = str(mig_path) # extract migration script id and name - mig_filename = mig_path_str.split('/')[-1].split('.')[0] + mig_filename = mig_path_str.split(os.sep)[-1].split('.')[0] name_fragments = mig_filename.split('_') mig_id = int(name_fragments[0]) diff --git a/manage/cmd/users.py b/manage/cmd/users.py index d4b477a..8887fa5 100644 --- a/manage/cmd/users.py +++ b/manage/cmd/users.py @@ -1,4 +1,7 @@ -from litecord.blueprints.auth import create_user +import base64 +import itsdangerous +import bcrypt +from litecord.blueprints.auth import create_user, make_token from litecord.enums import UserFlags @@ -9,7 +12,6 @@ async def find_user(username, discrim, ctx): WHERE username = $1 AND discriminator = $2 """, username, discrim) - async def set_user_staff(user_id, ctx): """Give a single user staff status.""" old_flags = await ctx.db.fetchval(""" @@ -52,6 +54,20 @@ async def make_staff(ctx, args): await set_user_staff(uid, ctx) print('OK: set staff') +async def generate_bot_token(ctx, args): + """Generate a token for specified bot.""" + + password_hash = await ctx.db.fetchval(""" + SELECT password_hash + FROM users + WHERE id = $1 AND bot = 'true' + """, int(args.user_id)) + + if not password_hash: + return print('cannot find a bot with specified id') + + print(make_token(args.user_id, password_hash)) + def setup(subparser): setup_test_parser = subparser.add_parser( @@ -82,3 +98,15 @@ def setup(subparser): ) staff_parser.set_defaults(func=make_staff) + + token_parser = subparser.add_parser( + 'generate_token', + help='generate a token for specified bot', + description=generate_bot_token.__doc__ + ) + + token_parser.add_argument( + 'user_id' + ) + + token_parser.set_defaults(func=generate_bot_token)