From 1087ed01868aaa17d03790af0a4350aa02075763 Mon Sep 17 00:00:00 2001 From: gabixdev Date: Mon, 3 Dec 2018 03:00:35 +0100 Subject: [PATCH 1/2] Manage command for retrieving user tokens --- manage/cmd/users.py | 37 ++++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/manage/cmd/users.py b/manage/cmd/users.py index d4b477a..92701d5 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,6 +12,13 @@ async def find_user(username, discrim, ctx): WHERE username = $1 AND discriminator = $2 """, username, discrim) +async def get_password_hash(id, ctx): + return await ctx.db.fetchval(""" + SELECT password_hash + FROM users + WHERE id = $1 + """, id) + async def set_user_staff(user_id, ctx): """Give a single user staff status.""" @@ -52,6 +62,16 @@ async def make_staff(ctx, args): await set_user_staff(uid, ctx) print('OK: set staff') +async def generate_token(ctx, args): + """Generate a token for specified user.""" + uid = await find_user(args.username, args.discrim, ctx) + + if not uid: + return print('user not found') + + password_hash = await get_password_hash(uid, ctx) + print(make_token(uid, password_hash)) + def setup(subparser): setup_test_parser = subparser.add_parser( @@ -82,3 +102,18 @@ def setup(subparser): ) staff_parser.set_defaults(func=make_staff) + + token_parser = subparser.add_parser( + 'generate_token', + help='generate a token for specified user', + description=generate_token.__doc__ + ) + + token_parser.add_argument( + 'username' + ) + token_parser.add_argument( + 'discrim', help='the discriminator of the user' + ) + + token_parser.set_defaults(func=generate_token) From 2542084c73186fd7c2362b0c7c652c5436a92031 Mon Sep 17 00:00:00 2001 From: gabixdev Date: Mon, 3 Dec 2018 03:26:50 +0100 Subject: [PATCH 2/2] Fix migration on windows + make generate_token generate only bot ids --- manage/cmd/migration/command.py | 8 +++---- manage/cmd/users.py | 39 ++++++++++++++------------------- 2 files changed, 20 insertions(+), 27 deletions(-) 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 92701d5..8887fa5 100644 --- a/manage/cmd/users.py +++ b/manage/cmd/users.py @@ -12,14 +12,6 @@ async def find_user(username, discrim, ctx): WHERE username = $1 AND discriminator = $2 """, username, discrim) -async def get_password_hash(id, ctx): - return await ctx.db.fetchval(""" - SELECT password_hash - FROM users - WHERE id = $1 - """, id) - - async def set_user_staff(user_id, ctx): """Give a single user staff status.""" old_flags = await ctx.db.fetchval(""" @@ -62,15 +54,19 @@ async def make_staff(ctx, args): await set_user_staff(uid, ctx) print('OK: set staff') -async def generate_token(ctx, args): - """Generate a token for specified user.""" - uid = await find_user(args.username, args.discrim, ctx) +async def generate_bot_token(ctx, args): + """Generate a token for specified bot.""" - if not uid: - return print('user not found') - - password_hash = await get_password_hash(uid, ctx) - print(make_token(uid, password_hash)) + 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): @@ -105,15 +101,12 @@ def setup(subparser): token_parser = subparser.add_parser( 'generate_token', - help='generate a token for specified user', - description=generate_token.__doc__ + help='generate a token for specified bot', + description=generate_bot_token.__doc__ ) token_parser.add_argument( - 'username' - ) - token_parser.add_argument( - 'discrim', help='the discriminator of the user' + 'user_id' ) - token_parser.set_defaults(func=generate_token) + token_parser.set_defaults(func=generate_bot_token)