diff --git a/images/.gitkeep b/images/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/litecord/blueprints/guilds.py b/litecord/blueprints/guilds.py index cd40ec2..2c80f34 100644 --- a/litecord/blueprints/guilds.py +++ b/litecord/blueprints/guilds.py @@ -86,6 +86,19 @@ async def guild_create_channels_prep(guild_id: int, channels: list): await create_guild_channel(guild_id, channel_id, ctype) +async def put_guild_icon(guild_id: int, icon: str): + """Insert a guild icon on the icon database.""" + if icon and icon.startswith('data'): + encoded = icon + else: + encoded = (f'data:image/jpeg;base64,{icon}' + if icon + else None) + + return await app.icons.put( + 'guild', guild_id, encoded, size=(128, 128)) + + @bp.route('', methods=['POST']) async def create_guild(): """Create a new guild, assigning @@ -96,13 +109,15 @@ async def create_guild(): guild_id = get_snowflake() + image = await put_guild_icon(guild_id, j['icon']) + await app.db.execute( """ INSERT INTO guilds (id, name, region, icon, owner_id, verification_level, default_message_notifications, explicit_content_filter) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - """, guild_id, j['name'], j['region'], j['icon'], user_id, + """, guild_id, j['name'], j['region'], image.icon_hash, user_id, j.get('verification_level', 0), j.get('default_message_notifications', 0), j.get('explicit_content_filter', 0)) @@ -157,8 +172,8 @@ async def get_guild(guild_id): ) -@bp.route('/', methods=['UPDATE']) -async def update_guild(guild_id): +@bp.route('/', methods=['PATCH']) +async def _update_guild(guild_id): user_id = await token_check() # TODO: check MANAGE_GUILD @@ -171,41 +186,62 @@ async def update_guild(guild_id): await app.db.execute(""" UPDATE guilds SET owner_id = $1 - WHERE guild_id = $2 + WHERE id = $2 """, int(j['owner_id']), guild_id) if 'name' in j: await app.db.execute(""" UPDATE guilds SET name = $1 - WHERE guild_id = $2 + WHERE id = $2 """, j['name'], guild_id) if 'region' in j: await app.db.execute(""" UPDATE guilds SET region = $1 - WHERE guild_id = $2 + WHERE id = $2 """, j['region'], guild_id) + if 'icon' in j: + # delete old + old_icon_hash = await app.db.fetchval(""" + SELECT icon + FROM guilds + WHERE id = $1 + """, guild_id) + + old_icon = await app.icons.get_guild_icon( + guild_id, old_icon_hash) + + await app.icons.delete(old_icon) + + new_icon = await put_guild_icon(guild_id, j['icon']) + + await app.db.execute(""" + UPDATE guilds + SET icon = $1 + WHERE id = $2 + """, new_icon.icon_hash, guild_id) + fields = ['verification_level', 'default_message_notifications', 'explicit_content_filter', 'afk_timeout'] for field in [f for f in fields if f in j]: - await app.db.execute(""" + await app.db.execute(f""" UPDATE guilds SET {field} = $1 - WHERE guild_id = $2 + WHERE id = $2 """, j[field], guild_id) channel_fields = ['afk_channel_id', 'system_channel_id'] for field in [f for f in channel_fields if f in j]: # TODO: check channel link to guild - await app.db.execute(""" + await app.db.execute(f""" UPDATE guilds SET {field} = $1 - WHERE guild_id = $2 + WHERE id = $2 """, j[field], guild_id) guild = await app.storage.get_guild_full( diff --git a/litecord/blueprints/icons.py b/litecord/blueprints/icons.py index cd09540..6d83159 100644 --- a/litecord/blueprints/icons.py +++ b/litecord/blueprints/icons.py @@ -1,33 +1,51 @@ +from os.path import splitext from quart import Blueprint, current_app as app, send_file bp = Blueprint('images', __name__) +async def send_icon(scope, key, icon_hash, **kwargs): + """Send an icon.""" + icon = await app.icons.generic_get( + scope, key, icon_hash, **kwargs) + + if not icon: + return '', 404 + + return await send_file(icon.as_path) + + +def splitext_(filepath): + name, ext = splitext(filepath) + return name, ext.strip('.') + + @bp.route('/emojis/.', methods=['GET']) -async def get_raw_emoji(emoji_id: int, ext: str): +async def _get_raw_emoji(emoji_id: int, ext: str): # emoji = app.icons.get_emoji(emoji_id, ext=ext) # just a test file for now return await send_file('./LICENSE') -@bp.route('/icons//.', methods=['GET']) -async def get_guild_icon(guild_id: int, icon_hash: str, ext: str): - pass +@bp.route('/icons//', methods=['GET']) +async def _get_guild_icon(guild_id: int, icon_file: str): + icon_hash, ext = splitext_(icon_file) + return await send_icon('guild', guild_id, icon_hash, ext=ext) @bp.route('/splashes//.', methods=['GET']) -async def get_guild_splash(guild_id: int, splash_hash: str, ext: str): +async def _get_guild_splash(guild_id: int, splash_hash: str, ext: str): pass @bp.route('/embed/avatars/.png') -async def get_default_user_avatar(discrim: int): +async def _get_default_user_avatar(discrim: int): pass @bp.route('/avatars//.') -async def get_user_avatar(user_id, avatar_hash, ext): - pass +async def _get_user_avatar(user_id, avatar_hash, ext): + return await send_icon('user', user_id, avatar_hash, ext=ext) # @bp.route('/app-icons//.') diff --git a/litecord/images.py b/litecord/images.py index c4d0978..e38c49e 100644 --- a/litecord/images.py +++ b/litecord/images.py @@ -1,6 +1,260 @@ +import mimetypes +import asyncio +import base64 + +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from io import BytesIO + +from logbook import Logger +from PIL import Image + +IMAGE_FOLDER = Path('./images') +log = Logger(__name__) + + +def _get_ext(mime: str): + extensions = mimetypes.guess_all_extensions(mime) + return extensions[0].strip('.') + + +def _get_mime(ext: str): + return mimetypes.types_map[f'.{ext}'] + + +@dataclass +class Icon: + """Main icon class""" + icon_hash: str + mime: str + + @property + def as_path(self) -> str: + """Return a filesystem path for the given icon.""" + ext = _get_ext(self.mime) + return str(IMAGE_FOLDER / f'{self.icon_hash}.{ext}') + + @property + def as_pathlib(self) -> str: + return Path(self.as_path) + + @property + def extension(self) -> str: + return _get_ext(self.mime) + + +class ImageError(Exception): + """Image error class.""" + pass + + +def to_raw(data_type: str, data: str) -> bytes: + """Given a data type in the data URI and data, + give the raw bytes being encoded.""" + if data_type == 'base64': + return base64.b64decode(data) + + return None + + +def _calculate_hash(fhandler) -> str: + """Generate a hash of the given file. + + This calls the seek(0) of the file handler + so it can be reused. + + Parameters + ---------- + fhandler: file object + Any file-like object. + + Returns + ------- + str + The SHA256 hash of the given file. + """ + hash_obj = sha256() + + for chunk in iter(lambda: fhandler.read(4096), b''): + hash_obj.update(chunk) + + # so that we can reuse the same handler + # later on + fhandler.seek(0) + + return hash_obj.hexdigest() + + +async def calculate_hash(fhandle, loop=None) -> str: + """Calculate a hash of the given file handle. + + Uses run_in_executor to do the job asynchronously so + the application doesn't lock up on large files. + """ + if not loop: + loop = asyncio.get_event_loop() + + fut = loop.run_in_executor(None, _calculate_hash, fhandle) + return await fut + + +def parse_data_uri(string) -> tuple: + """Extract image data.""" + try: + header, headered_data = string.split(';') + + _, given_mime = header.split(':') + data_type, data = headered_data.split(',') + + raw_data = to_raw(data_type, data) + if raw_data is None: + raise ImageError('Unknown data header') + + return given_mime, raw_data + except ValueError: + raise ImageError('data URI invalid syntax') -MIMES = { - } class IconManager: """Main icon manager.""" + def __init__(self, app): + self.app = app + self.storage = app.storage + + async def _convert_ext(self, icon: Icon, target: str): + target_mime = _get_mime(target) + log.info('converting from {} to {}', icon.mime, target_mime) + + target_path = IMAGE_FOLDER / f'{icon.icon_hash}.{target}' + + if target_path.exists(): + return Icon(icon.icon_hash, target_mime) + + image = Image.open(icon.as_path) + target_fd = target_path.open('wb') + image.save(target_fd, format=target) + target_fd.close() + + return Icon(icon.icon_hash, target_mime) + + async def generic_get(self, scope, key, icon_hash, **kwargs) -> Icon: + """Get any icon.""" + key = str(key) + + icon_row = await self.storage.db.fetchrow(""" + SELECT hash, mime + FROM icons + WHERE scope = $1 + AND key = $2 + AND hash = $3 + """, scope, key, icon_hash) + + if not icon_row: + return None + + icon = Icon(icon_row['hash'], icon_row['mime']) + + if not icon.as_pathlib.exists(): + await self.delete(icon) + return None + + if 'ext' in kwargs and kwargs['ext'] != icon.extension: + return await self._convert_ext(icon, kwargs['ext']) + + return icon + + async def get_guild_icon(self, guild_id: int, icon_hash: str, **kwargs): + """Get an icon for a guild.""" + return await self.generic_get( + 'guild', guild_id, icon_hash, **kwargs) + + async def put(self, scope: str, key: str, + b64_data: str, **kwargs) -> Icon: + """Insert an icon.""" + if b64_data is None: + return Icon(None, '') + + mime, raw_data = parse_data_uri(b64_data) + data_fd = BytesIO(raw_data) + + # get an extension for the given data uri + extension = _get_ext(mime) + + if 'size' in kwargs: + image = Image.open(data_fd) + + want = kwargs['size'] + + log.info('resizing from {} to {}', + image.size, want) + + resized = image.resize(want) + + data_fd = BytesIO() + resized.save(data_fd, format=extension) + + # reseek to copy it to raw_data + data_fd.seek(0) + raw_data = data_fd.read() + + data_fd.seek(0) + + # calculate sha256 + icon_hash = await calculate_hash(data_fd) + + await self.storage.db.execute(""" + INSERT INTO icons (scope, key, hash, mime) + VALUES ($1, $2, $3, $4) + """, scope, str(key), icon_hash, mime) + + # write it off to fs + icon_path = IMAGE_FOLDER / f'{icon_hash}.{extension}' + icon_path.write_bytes(raw_data) + + # copy from data_fd to icon_fd + # with icon_path.open(mode='wb') as icon_fd: + # icon_fd.write(data_fd.read()) + + return Icon(icon_hash, mime) + + async def delete(self, icon: Icon): + """Delete an icon from the database and filesystem.""" + if not icon: + return + + # dereference + await self.storage.db.execute(""" + UPDATE users + SET avatar = NULL + WHERE avatar = $1 + """, icon.icon_hash) + + await self.storage.db.execute(""" + UPDATE group_dm_channels + SET icon = NULL + WHERE icon = $1 + """, icon.icon_hash) + + await self.storage.db.execute(""" + DELETE FROM guild_emoji + WHERE image = $1 + """, icon.icon_hash) + + await self.storage.db.execute(""" + UPDATE guilds + SET icon = NULL + WHERE icon = $1 + """, icon.icon_hash) + + await self.storage.db.execute(""" + DELETE FROM icons + WHERE hash = $1 + """, icon.icon_hash) + + icon_path = icon.as_pathlib + + try: + icon_path.unlink() + except FileNotFoundError: + pass diff --git a/litecord/schemas.py b/litecord/schemas.py index 346c60b..a234c41 100644 --- a/litecord/schemas.py +++ b/litecord/schemas.py @@ -234,20 +234,22 @@ GUILD_UPDATE = { }, 'region': {'type': 'voice_region', 'required': False}, 'icon': {'type': 'b64_icon', 'required': False}, + 'splash': {'type': 'b64_icon', 'required': False, 'nullable': True}, - 'verification_level': {'type': 'verification_level', 'required': False}, + 'verification_level': { + 'type': 'verification_level', 'required': False}, 'default_message_notifications': { - 'type': 'msg_notifications', - 'required': False, - }, + 'type': 'msg_notifications', 'required': False}, 'explicit_content_filter': {'type': 'explicit', 'required': False}, - 'afk_channel_id': {'type': 'snowflake', 'required': False}, + 'afk_channel_id': { + 'type': 'snowflake', 'required': False, 'nullable': True}, 'afk_timeout': {'type': 'number', 'required': False}, 'owner_id': {'type': 'snowflake', 'required': False}, - 'system_channel_id': {'type': 'snowflake', 'required': False}, + 'system_channel_id': { + 'type': 'snowflake', 'required': False, 'nullable': True}, } @@ -451,6 +453,7 @@ INVITE = { 'temporary': {'type': 'boolean', 'required': False, 'default': False}, 'unique': {'type': 'boolean', 'required': False, 'default': True}, + 'validate': {'type': 'boolean', 'required': False, 'nullable': True} } USER_SETTINGS = { diff --git a/manage/cmd/migration/scripts/2_icons_table.sql b/manage/cmd/migration/scripts/2_icons_table.sql new file mode 100644 index 0000000..ffe6fd9 --- /dev/null +++ b/manage/cmd/migration/scripts/2_icons_table.sql @@ -0,0 +1,35 @@ + +-- new icons table +CREATE TABLE IF NOT EXISTS icons ( + scope text NOT NULL, + key text, + hash text UNIQUE NOT NULL, + mime text NOT NULL, + PRIMARY KEY (scope, hash, mime) +); + +-- dummy attachments table for now. +CREATE TABLE IF NOT EXISTS attachments ( + id bigint NOT NULL, + PRIMARY KEY (id) +); + +-- remove the old columns referencing the files table +ALTER TABLE users DROP COLUMN avatar; +ALTER TABLE users ADD COLUMN avatar text REFERENCES icons (hash) DEFAULT NULL; + +ALTER TABLE group_dm_channels DROP COLUMN icon; +ALTER TABLE group_dm_channels ADD COLUMN icon text REFERENCES icons (hash); + +ALTER TABLE guild_emoji DROP COLUMN image; +ALTER TABLE guild_emoji ADD COLUMN image text REFERENCES icons (hash); + +ALTER TABLE guilds DROP COLUMN icon; +ALTER TABLE guilds ADD COLUMN icon text REFERENCES icons (hash) DEFAULT NULL; + +-- this one is a change from files to the attachments table +ALTER TABLE message_attachments DROP COLUMN attachment; +ALTER TABLE guild_emoji ADD COLUMN attachment bigint REFERENCES attachments (id); + +-- remove files table +DROP TABLE files; diff --git a/run.py b/run.py index 8699cf4..56d2be0 100644 --- a/run.py +++ b/run.py @@ -37,6 +37,7 @@ from litecord.gateway.state_manager import StateManager from litecord.storage import Storage from litecord.dispatcher import EventDispatcher from litecord.presence import PresenceManager +from litecord.images import IconManager # setup logbook handler = StreamHandler(sys.stdout, level=logbook.INFO) @@ -91,7 +92,11 @@ def set_blueprints(app_): } for bp, suffix in bps.items(): - url_prefix = f'/api/v6/{suffix or ""}' if suffix != -1 else '' + url_prefix = f'/api/v6{suffix or ""}' + + if suffix == -1: + url_prefix = '' + app_.register_blueprint(bp, url_prefix=url_prefix) @@ -163,6 +168,7 @@ def init_app_managers(app): app.ratelimiter = RatelimitManager() app.state_manager = StateManager() app.storage = Storage(app.db) + app.icons = IconManager(app) app.dispatcher = EventDispatcher(app) app.presence = PresenceManager(app.storage, diff --git a/schema.sql b/schema.sql index 7f66a7b..dd1b56c 100644 --- a/schema.sql +++ b/schema.sql @@ -40,19 +40,24 @@ INSERT INTO user_conn_apps (id, name) VALUES (9, 'Skype'); INSERT INTO user_conn_apps (id, name) VALUES (10, 'League of Legends'); -CREATE TABLE IF NOT EXISTS files ( - -- snowflake id of the file - id bigint PRIMARY KEY NOT NULL, +CREATE TABLE IF NOT EXISTS icons ( + -- can be 'user', 'guild', 'emoji' + scope text NOT NULL, - -- sha512(file) - hash text NOT NULL, - mimetype text NOT NULL, + -- can be a user snowflake, guild snowflake or + -- emoji snowflake + key text, - -- path to the file system - fspath text NOT NULL + -- sha256 of the icon + hash text UNIQUE NOT NULL, + + -- icon mime + mime text NOT NULL, + PRIMARY KEY (scope, hash, mime) ); + CREATE TABLE IF NOT EXISTS users ( id bigint UNIQUE NOT NULL, username text NOT NULL, @@ -63,7 +68,7 @@ CREATE TABLE IF NOT EXISTS users ( bot boolean DEFAULT FALSE, mfa_enabled boolean DEFAULT FALSE, verified boolean DEFAULT FALSE, - avatar bigint REFERENCES files (id) DEFAULT NULL, + avatar text REFERENCES icons (hash) DEFAULT NULL, -- user badges, discord dev, etc flags int DEFAULT 0, @@ -320,7 +325,7 @@ CREATE TABLE IF NOT EXISTS group_dm_channels ( id bigint REFERENCES channels (id) ON DELETE CASCADE, owner_id bigint REFERENCES users (id), name text, - icon bigint REFERENCES files (id), + icon text REFERENCES icons (hash) DEFAULT NULL, PRIMARY KEY (id) ); @@ -359,7 +364,7 @@ CREATE TABLE IF NOT EXISTS guild_emoji ( uploader_id bigint REFERENCES users (id), name text NOT NULL, - image bigint REFERENCES files (id), + image text REFERENCES icons (hash), animated bool DEFAULT false, managed bool DEFAULT false, require_colons bool DEFAULT false @@ -521,7 +526,7 @@ CREATE TABLE IF NOT EXISTS messages ( CREATE TABLE IF NOT EXISTS message_attachments ( message_id bigint REFERENCES messages (id), - attachment bigint REFERENCES files (id), + attachment bigint REFERENCES attachments (id), PRIMARY KEY (message_id, attachment) );