all: add guild icon support

- fix update_guild's methods
 - fix update_guild's sql statements
 - litecord: add images module
 - schemas: add splash to GUILD_UPDATE
 - schemas: add validate to INVITE
 - manage.cmd.migration: add script 2

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('/<int:guild_id>', methods=['UPDATE'])
+@bp.route('/<int:guild_id>', methods=['PATCH'])
-async def update_guild(guild_id):
+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(

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/<int:emoji_id>.<ext>', 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/<int:guild_id>/<icon_hash>.<ext>', methods=['GET'])
+@bp.route('/icons/<int:guild_id>/<icon_file>', methods=['GET'])
-async def get_guild_icon(guild_id: int, icon_hash: str, ext: str):
+async def _get_guild_icon(guild_id: int, icon_file: str):
-    pass
+    icon_hash, ext = splitext_(icon_file)
+    return await send_icon('guild', guild_id, icon_hash, ext=ext)
 
 
 @bp.route('/splashes/<int:guild_id>/<icon_hash>.<ext>', 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/<int:discrim>.png')
-async def get_default_user_avatar(discrim: int):
+async def _get_default_user_avatar(discrim: int):
     pass
 
 
 @bp.route('/avatars/<int:user_id>/<avatar_hash>.<ext>')
-async def get_user_avatar(user_id, avatar_hash, ext):
+async def _get_user_avatar(user_id, avatar_hash, ext):
-    pass
+    return await send_icon('user', user_id, avatar_hash, ext=ext)
 
 
 # @bp.route('/app-icons/<int:application_id>/<icon_hash>.<ext>')

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

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',
+        'type': 'msg_notifications', 'required': False},
-        '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 = {

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;

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,

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 (
+CREATE TABLE IF NOT EXISTS icons (
-    -- snowflake id of the file
+    -- can be 'user', 'guild', 'emoji'
-    id bigint PRIMARY KEY NOT NULL,
+    scope text NOT NULL,
 
-    -- sha512(file)
+    -- can be a user snowflake, guild snowflake or
-    hash text NOT NULL,
+    -- emoji snowflake
-    mimetype text NOT NULL,
+    key text,
 
-    -- path to the file system
+    -- sha256 of the icon
-    fspath text NOT NULL
+    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)
 );