diff --git a/litecord/auth.py b/litecord/auth.py
index b96e475..0341fb9 100644
--- a/litecord/auth.py
+++ b/litecord/auth.py
@@ -20,15 +20,14 @@ along with this program. If not, see .
import base64
import binascii
from random import randint
-from typing import Tuple, Optional
+from typing import Tuple
import bcrypt
-from asyncpg import UniqueViolationError
from itsdangerous import TimestampSigner, BadSignature
from logbook import Logger
from quart import request, current_app as app
-from litecord.errors import Forbidden, Unauthorized, BadRequest
+from litecord.errors import Forbidden, Unauthorized
from litecord.snowflake import get_snowflake
from litecord.enums import UserFlags
@@ -150,105 +149,3 @@ async def hash_data(data: str, loop=None) -> str:
hashed = await loop.run_in_executor(None, bcrypt.hashpw, buf, bcrypt.gensalt(14))
return hashed.decode()
-
-
-async def check_username_usage(username: str):
- """Raise an error if too many people are with the same username."""
- same_username = await app.db.fetchval(
- """
- SELECT COUNT(*)
- FROM users
- WHERE username = $1
- """,
- username,
- )
-
- if same_username > 9000:
- raise BadRequest(
- "Too many people.",
- {
- "username": "Too many people used the same username. "
- "Please choose another"
- },
- )
-
-
-def _raw_discrim() -> str:
- discrim_number = randint(1, 9999)
- return "%04d" % discrim_number
-
-
-async def roll_discrim(username: str) -> Optional[str]:
- """Roll a discriminator for a DiscordTag.
-
- Tries to generate one 10 times.
-
- Calls check_username_usage.
- """
-
- # we shouldn't roll discrims for usernames
- # that have been used too much.
- await check_username_usage(username)
-
- # max 10 times for a reroll
- for _ in range(10):
- # generate random discrim
- discrim = _raw_discrim()
-
- # check if anyone is with it
- res = await app.db.fetchval(
- """
- SELECT id
- FROM users
- WHERE username = $1 AND discriminator = $2
- """,
- username,
- discrim,
- )
-
- # if no user is found with the (username, discrim)
- # pair, then this is unique! return it.
- if res is None:
- return discrim
-
- return None
-
-
-async def create_user(username: str, email: str, password: str) -> Tuple[int, str]:
- """Create a single user.
-
- Generates a distriminator and other information. You can fetch the user
- data back with :meth:`Storage.get_user`.
- """
- db = app.db
- loop = app.loop
-
- new_id = get_snowflake()
- new_discrim = await roll_discrim(username)
-
- if new_discrim is None:
- raise BadRequest(
- "Unable to register.",
- {"username": "Too many people are with this username."},
- )
-
- pwd_hash = await hash_data(password, loop)
-
- try:
- await db.execute(
- """
- INSERT INTO users
- (id, email, username, discriminator, password_hash)
- VALUES
- ($1, $2, $3, $4, $5)
- """,
- new_id,
- email,
- username,
- new_discrim,
- pwd_hash,
- )
- except UniqueViolationError:
- raise BadRequest("Email already used.")
-
- return new_id, pwd_hash
diff --git a/litecord/blueprints/users.py b/litecord/blueprints/users.py
index 667e143..7009bd5 100644
--- a/litecord/blueprints/users.py
+++ b/litecord/blueprints/users.py
@@ -36,46 +36,12 @@ from litecord.permissions import base_permissions
from litecord.blueprints.auth import check_password
from litecord.utils import to_update
+from litecord.common.users import mass_user_update, delete_user
bp = Blueprint("user", __name__)
log = Logger(__name__)
-async def mass_user_update(user_id):
- """Dispatch USER_UPDATE in a mass way."""
- # by using dispatch_with_filter
- # we're guaranteeing all shards will get
- # a USER_UPDATE once and not any others.
-
- session_ids = []
-
- public_user = await app.storage.get_user(user_id)
- private_user = await app.storage.get_user(user_id, secure=True)
-
- session_ids.extend(
- await app.dispatcher.dispatch_user(user_id, "USER_UPDATE", private_user)
- )
-
- guild_ids = await app.user_storage.get_user_guilds(user_id)
- friend_ids = await app.user_storage.get_friend_ids(user_id)
-
- session_ids.extend(
- await app.dispatcher.dispatch_many_filter_list(
- "guild", guild_ids, session_ids, "USER_UPDATE", public_user
- )
- )
-
- session_ids.extend(
- await app.dispatcher.dispatch_many_filter_list(
- "friend", friend_ids, session_ids, "USER_UPDATE", public_user
- )
- )
-
- await app.dispatcher.dispatch_many("lazy_guild", guild_ids, "update_user", user_id)
-
- return public_user, private_user
-
-
@bp.route("/@me", methods=["GET"])
async def get_me():
"""Get the current user's information."""
@@ -472,113 +438,6 @@ def rand_hex(length: int = 8) -> str:
return urandom(length).hex()[:length]
-async def _del_from_table(db, table: str, user_id: int):
- """Delete a row from a table."""
- column = {
- "channel_overwrites": "target_user",
- "user_settings": "id",
- "group_dm_members": "member_id",
- }.get(table, "user_id")
-
- res = await db.execute(
- f"""
- DELETE FROM {table}
- WHERE {column} = $1
- """,
- user_id,
- )
-
- log.info("Deleting uid {} from {}, res: {!r}", user_id, table, res)
-
-
-async def delete_user(user_id, *, mass_update: bool = True):
- """Delete a user. Does not disconnect the user."""
- db = app.db
-
- new_username = f"Deleted User {rand_hex()}"
-
- # by using a random hex in password_hash
- # we break attempts at using the default '123' password hash
- # to issue valid tokens for deleted users.
-
- await db.execute(
- """
- UPDATE users
- SET
- username = $1,
- email = NULL,
- mfa_enabled = false,
- verified = false,
- avatar = NULL,
- flags = 0,
- premium_since = NULL,
- phone = '',
- password_hash = $2
- WHERE
- id = $3
- """,
- new_username,
- rand_hex(32),
- user_id,
- )
-
- # remove the user from various tables
- await _del_from_table(db, "user_settings", user_id)
- await _del_from_table(db, "user_payment_sources", user_id)
- await _del_from_table(db, "user_subscriptions", user_id)
- await _del_from_table(db, "user_payments", user_id)
- await _del_from_table(db, "user_read_state", user_id)
- await _del_from_table(db, "guild_settings", user_id)
- await _del_from_table(db, "guild_settings_channel_overrides", user_id)
-
- await db.execute(
- """
- DELETE FROM relationships
- WHERE user_id = $1 OR peer_id = $1
- """,
- user_id,
- )
-
- # DMs are still maintained, but not the state.
- await _del_from_table(db, "dm_channel_state", user_id)
-
- # NOTE: we don't delete the group dms the user is an owner of...
- # TODO: group dm owner reassign when the owner leaves a gdm
- await _del_from_table(db, "group_dm_members", user_id)
-
- await _del_from_table(db, "members", user_id)
- await _del_from_table(db, "member_roles", user_id)
- await _del_from_table(db, "channel_overwrites", user_id)
-
- # after updating the user, we send USER_UPDATE so that all the other
- # clients can refresh their caches on the now-deleted user
- if mass_update:
- await mass_user_update(user_id)
-
-
-async def user_disconnect(user_id: int):
- """Disconnects the given user's devices."""
- # after removing the user from all tables, we need to force
- # all known user states to reconnect, causing the user to not
- # be online anymore.
- user_states = app.state_manager.user_states(user_id)
-
- for state in user_states:
- # make it unable to resume
- app.state_manager.remove(state)
-
- if not state.ws:
- continue
-
- # force a close, 4000 should make the client reconnect.
- await state.ws.ws.close(4000)
-
- # force everyone to see the user as offline
- await app.presence.dispatch_pres(
- user_id, {"afk": False, "status": "offline", "game": None, "since": 0}
- )
-
-
@bp.route("/@me/delete", methods=["POST"])
async def delete_account():
"""Delete own account.
diff --git a/litecord/common/users.py b/litecord/common/users.py
new file mode 100644
index 0000000..bd67dd8
--- /dev/null
+++ b/litecord/common/users.py
@@ -0,0 +1,250 @@
+from random import randint
+from typing import Tuple, Optional
+
+from quart import current_app as app
+from asyncpg import UniqueViolationError
+
+from ..snowflake import get_snowflake
+from ..errors import BadRequest
+from ..auth import hash_data
+
+
+async def mass_user_update(user_id):
+ """Dispatch USER_UPDATE in a mass way."""
+ # by using dispatch_with_filter
+ # we're guaranteeing all shards will get
+ # a USER_UPDATE once and not any others.
+
+ session_ids = []
+
+ public_user = await app.storage.get_user(user_id)
+ private_user = await app.storage.get_user(user_id, secure=True)
+
+ session_ids.extend(
+ await app.dispatcher.dispatch_user(user_id, "USER_UPDATE", private_user)
+ )
+
+ guild_ids = await app.user_storage.get_user_guilds(user_id)
+ friend_ids = await app.user_storage.get_friend_ids(user_id)
+
+ session_ids.extend(
+ await app.dispatcher.dispatch_many_filter_list(
+ "guild", guild_ids, session_ids, "USER_UPDATE", public_user
+ )
+ )
+
+ session_ids.extend(
+ await app.dispatcher.dispatch_many_filter_list(
+ "friend", friend_ids, session_ids, "USER_UPDATE", public_user
+ )
+ )
+
+ await app.dispatcher.dispatch_many("lazy_guild", guild_ids, "update_user", user_id)
+
+ return public_user, private_user
+
+
+async def check_username_usage(username: str):
+ """Raise an error if too many people are with the same username."""
+ same_username = await app.db.fetchval(
+ """
+ SELECT COUNT(*)
+ FROM users
+ WHERE username = $1
+ """,
+ username,
+ )
+
+ if same_username > 9000:
+ raise BadRequest(
+ "Too many people.",
+ {
+ "username": "Too many people used the same username. "
+ "Please choose another"
+ },
+ )
+
+
+def _raw_discrim() -> str:
+ discrim_number = randint(1, 9999)
+ return "%04d" % discrim_number
+
+
+async def roll_discrim(username: str) -> Optional[str]:
+ """Roll a discriminator for a DiscordTag.
+
+ Tries to generate one 10 times.
+
+ Calls check_username_usage.
+ """
+
+ # we shouldn't roll discrims for usernames
+ # that have been used too much.
+ await check_username_usage(username)
+
+ # max 10 times for a reroll
+ for _ in range(10):
+ # generate random discrim
+ discrim = _raw_discrim()
+
+ # check if anyone is with it
+ res = await app.db.fetchval(
+ """
+ SELECT id
+ FROM users
+ WHERE username = $1 AND discriminator = $2
+ """,
+ username,
+ discrim,
+ )
+
+ # if no user is found with the (username, discrim)
+ # pair, then this is unique! return it.
+ if res is None:
+ return discrim
+
+ return None
+
+
+async def create_user(username: str, email: str, password: str) -> Tuple[int, str]:
+ """Create a single user.
+
+ Generates a distriminator and other information. You can fetch the user
+ data back with :meth:`Storage.get_user`.
+ """
+ new_id = get_snowflake()
+ new_discrim = await roll_discrim(username)
+
+ if new_discrim is None:
+ raise BadRequest(
+ "Unable to register.",
+ {"username": "Too many people are with this username."},
+ )
+
+ pwd_hash = await hash_data(password)
+
+ try:
+ await app.db.execute(
+ """
+ INSERT INTO users
+ (id, email, username, discriminator, password_hash)
+ VALUES
+ ($1, $2, $3, $4, $5)
+ """,
+ new_id,
+ email,
+ username,
+ new_discrim,
+ pwd_hash,
+ )
+ except UniqueViolationError:
+ raise BadRequest("Email already used.")
+
+ return new_id, pwd_hash
+
+
+async def _del_from_table(db, table: str, user_id: int):
+ """Delete a row from a table."""
+ column = {
+ "channel_overwrites": "target_user",
+ "user_settings": "id",
+ "group_dm_members": "member_id",
+ }.get(table, "user_id")
+
+ res = await db.execute(
+ f"""
+ DELETE FROM {table}
+ WHERE {column} = $1
+ """,
+ user_id,
+ )
+
+ log.info("Deleting uid {} from {}, res: {!r}", user_id, table, res)
+
+
+async def delete_user(user_id, *, mass_update: bool = True):
+ """Delete a user. Does not disconnect the user."""
+ db = app.db
+
+ new_username = f"Deleted User {rand_hex()}"
+
+ # by using a random hex in password_hash
+ # we break attempts at using the default '123' password hash
+ # to issue valid tokens for deleted users.
+
+ await db.execute(
+ """
+ UPDATE users
+ SET
+ username = $1,
+ email = NULL,
+ mfa_enabled = false,
+ verified = false,
+ avatar = NULL,
+ flags = 0,
+ premium_since = NULL,
+ phone = '',
+ password_hash = $2
+ WHERE
+ id = $3
+ """,
+ new_username,
+ rand_hex(32),
+ user_id,
+ )
+
+ # remove the user from various tables
+ await _del_from_table(db, "user_settings", user_id)
+ await _del_from_table(db, "user_payment_sources", user_id)
+ await _del_from_table(db, "user_subscriptions", user_id)
+ await _del_from_table(db, "user_payments", user_id)
+ await _del_from_table(db, "user_read_state", user_id)
+ await _del_from_table(db, "guild_settings", user_id)
+ await _del_from_table(db, "guild_settings_channel_overrides", user_id)
+
+ await db.execute(
+ """
+ DELETE FROM relationships
+ WHERE user_id = $1 OR peer_id = $1
+ """,
+ user_id,
+ )
+
+ # DMs are still maintained, but not the state.
+ await _del_from_table(db, "dm_channel_state", user_id)
+
+ # NOTE: we don't delete the group dms the user is an owner of...
+ # TODO: group dm owner reassign when the owner leaves a gdm
+ await _del_from_table(db, "group_dm_members", user_id)
+
+ await _del_from_table(db, "members", user_id)
+ await _del_from_table(db, "member_roles", user_id)
+ await _del_from_table(db, "channel_overwrites", user_id)
+
+ # after updating the user, we send USER_UPDATE so that all the other
+ # clients can refresh their caches on the now-deleted user
+ if mass_update:
+ await mass_user_update(user_id)
+
+
+async def user_disconnect(user_id: int):
+ """Disconnects the given user's devices."""
+ # after removing the user from all tables, we need to force
+ # all known user states to reconnect, causing the user to not
+ # be online anymore.
+ user_states = app.state_manager.user_states(user_id)
+
+ for state in user_states:
+ # make it unable to resume
+ app.state_manager.remove(state)
+
+ if not state.ws:
+ continue
+
+ # force a close, 4000 should make the client reconnect.
+ await state.ws.ws.close(4000)
+
+ # force everyone to see the user as offline
+ await app.presence.dispatch_pres(
+ user_id, {"afk": False, "status": "offline", "game": None, "since": 0}
+ )
diff --git a/manage/cmd/users.py b/manage/cmd/users.py
index c20e181..50818e8 100644
--- a/manage/cmd/users.py
+++ b/manage/cmd/users.py
@@ -17,9 +17,8 @@ along with this program. If not, see .
"""
-from litecord.auth import create_user
+from litecord.common.users import create_user, delete_user
from litecord.blueprints.auth import make_token
-from litecord.blueprints.users import delete_user
from litecord.enums import UserFlags