Merge branch 'message-attachments' into 'master'

Message attachments

Closes #5

See merge request litecord/litecord!14

.gitignore

@@ -105,5 +105,6 @@ venv.bak/
 
 config.py
 images/*
+attachments/*
 
 .DS_Store

litecord/blueprints/__init__.py

@@ -31,6 +31,8 @@ from .dms import bp as dms
 from .icons import bp as icons
 from .nodeinfo import bp as nodeinfo
 from .static import bp as static
+from .attachments import bp as attachments
+
 __all__ = ['gateway', 'auth', 'users', 'guilds', 'channels',
            'webhooks', 'science', 'voice', 'invites', 'relationships',
-           'dms', 'icons', 'nodeinfo', 'static']
+           'dms', 'icons', 'nodeinfo', 'static', 'attachments']

litecord/blueprints/attachments.py

@@ -0,0 +1,117 @@
+"""
+
+Litecord
+Copyright (C) 2018  Luna Mendes
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, version 3 of the License.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""
+
+from pathlib import Path
+
+from quart import Blueprint, send_file, current_app as app, request
+from PIL import Image
+
+from litecord.images import resize_gif
+
+bp = Blueprint('attachments', __name__)
+ATTACHMENTS = Path.cwd() / 'attachments'
+
+
+async def _resize_gif(attach_id: int, resized_path: Path,
+                      width: int, height: int) -> str:
+    """Resize a GIF attachment."""
+
+    # get original gif bytes
+    orig_path = ATTACHMENTS / f'{attach_id}.gif'
+    orig_bytes = orig_path.read_bytes()
+
+    # give them and the target size to the
+    # image module's resize_gif
+
+    _data_fd, raw_data = await resize_gif(orig_bytes, (width, height))
+
+    # write raw_data to the destination
+    resized_path.write_bytes(raw_data)
+
+    return str(resized_path)
+
+
+async def _resize(image, attach_id: int, ext: str,
+                  width: int, height: int) -> str:
+    """Resize an image."""
+    # check if we have it on the folder
+    resized_path = ATTACHMENTS / f'{attach_id}_{width}_{height}.{ext}'
+
+    # keep a str-fied instance since that is what
+    # we'll return.
+    resized_path_s = str(resized_path)
+
+    if resized_path.exists():
+        return resized_path_s
+
+    # if we dont, we need to generate it off the
+    # given image instance.
+
+    # the process is different for gif files because we need
+    # gifsicle. doing it manually is too troublesome.
+    if ext == 'gif':
+        return await _resize_gif(attach_id, resized_path, width, height)
+
+    # NOTE: this is the same resize mode for icons.
+    resized = image.resize((width, height), resample=Image.LANCZOS)
+    resized.save(resized_path_s, format=ext)
+
+    return resized_path_s
+
+
+@bp.route('/attachments'
+          '/<int:channel_id>/<int:message_id>/<filename>',
+          methods=['GET'])
+async def _get_attachment(channel_id: int, message_id: int,
+                          filename: str):
+
+    attach_id = await app.db.fetchval("""
+    SELECT id
+    FROM attachments
+    WHERE channel_id = $1
+      AND message_id = $2
+      AND filename = $3
+    """, channel_id, message_id, filename)
+
+    if attach_id is None:
+        return '', 404
+
+    ext = filename.split('.')[-1]
+    filepath = f'./attachments/{attach_id}.{ext}'
+
+    image = Image.open(filepath)
+    im_width, im_height = image.size
+
+    try:
+        width = int(request.args.get('width', 0)) or im_width
+    except ValueError:
+        return '', 400
+
+    try:
+        height = int(request.args.get('height', 0)) or im_height
+    except ValueError:
+        return '', 400
+
+    # if width and height are the same (happens if they weren't provided)
+    if width == im_width and height == im_height:
+        return await send_file(filepath)
+
+    # resize image
+    new_filepath = await _resize(image, attach_id, ext, width, height)
+    return await send_file(new_filepath)

litecord/blueprints/channel/messages.py

@@ -18,11 +18,12 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
 
 import re
+import json
 
+from PIL import Image
 from quart import Blueprint, request, current_app as app, jsonify
 from logbook import Logger
 
-
 from litecord.blueprints.auth import token_check
 from litecord.blueprints.checks import channel_check, channel_perm_check
 from litecord.blueprints.dms import try_dm_state
@@ -311,9 +312,105 @@ async def process_url_embed(config, storage, dispatcher, session, payload: dict)
         'channel', channel_id, 'MESSAGE_UPDATE', update_payload)
 
 
+async def _msg_input() -> tuple:
+    """Extract the json input and any file information
+    the client gave to us in the request.
+
+    This only applies to create message route.
+    """
+    form = await request.form
+    request_json = await request.get_json() or {}
+
+    # NOTE: embed isn't set on form data
+    json_from_form = {
+        'content': form.get('content', ''),
+        'nonce': form.get('nonce', '0'),
+        'tts': json.loads(form.get('tts', 'false')),
+    }
+
+    json_from_form.update(request_json)
+
+    files = await request.files
+    # we don't really care about the given fields on the files dict
+    return json_from_form, [v for k, v in files.items()]
+
+
+def _check_content(payload: dict, files: list):
+    """Check if there is actually any content being sent to us."""
+    has_content = bool(payload.get('content', ''))
+    has_embed = 'embed' in payload
+    has_files = len(files) > 0
+
+    has_total_content = has_content or has_embed or has_files
+
+    if not has_total_content:
+        raise BadRequest('No content has been provided.')
+
+
+async def _add_attachment(message_id: int, channel_id: int,
+                          attachment_file) -> int:
+    """Add an attachment to a message.
+
+    Parameters
+    ----------
+    message_id: int
+        The ID of the message getting the attachment.
+    attachment_file: quart.FileStorage
+        quart FileStorage instance of the file.
+    """
+
+    attachment_id = get_snowflake()
+    filename = attachment_file.filename
+
+    # understand file info
+    mime = attachment_file.mimetype
+    is_image = mime.startswith('image/')
+
+    img_width, img_height = None, None
+
+    # extract file size
+    # TODO: this is probably inneficient
+    file_size = attachment_file.stream.getbuffer().nbytes
+
+    if is_image:
+        # open with pillow, extract image size
+        image = Image.open(attachment_file.stream)
+        img_width, img_height = image.size
+
+        # NOTE: DO NOT close the image, as closing the image will
+        # also close the stream.
+
+        # reset it to 0 for later usage
+        attachment_file.stream.seek(0)
+
+    await app.db.execute(
+        """
+        INSERT INTO attachments
+            (id, channel_id, message_id,
+             filename, filesize,
+             image, width, height)
+        VALUES
+            ($1, $2, $3, $4, $5, $6, $7, $8)
+        """,
+        attachment_id, channel_id, message_id,
+        filename, file_size,
+        is_image, img_width, img_height)
+
+    ext = filename.split('.')[-1]
+
+    with open(f'attachments/{attachment_id}.{ext}', 'wb') as attach_file:
+        attach_file.write(attachment_file.stream.read())
+
+    log.debug('written {} bytes for attachment id {}',
+              file_size, attachment_id)
+
+    return attachment_id
+
+
 @bp.route('/<int:channel_id>/messages', methods=['POST'])
 async def _create_message(channel_id):
     """Create a message."""
+
     user_id = await token_check()
     ctype, guild_id = await channel_check(user_id, channel_id)
 
@@ -323,7 +420,11 @@ async def _create_message(channel_id):
         await channel_perm_check(user_id, channel_id, 'send_messages')
         actual_guild_id = guild_id
 
-    j = validate(await request.get_json(), MESSAGE_CREATE)
+    payload_json, files = await _msg_input()
+    j = validate(payload_json, MESSAGE_CREATE)
+
+    print(payload_json, files)
+    _check_content(payload_json, files)
 
     # TODO: check connection to the gateway
 
@@ -355,6 +456,10 @@ async def _create_message(channel_id):
                        if 'embed' in j else []),
         })
 
+    # for each file given, we add it as an attachment
+    for pre_attachment in files:
+        await _add_attachment(message_id, channel_id, pre_attachment)
+
     payload = await app.storage.get_message(message_id, user_id)
 
     if ctype == ChannelType.DM:

litecord/images.py

@@ -49,7 +49,7 @@ MIMES = {
 }
 
 
-def _get_ext(mime: str) -> str:
+def get_ext(mime: str) -> str:
     if mime in EXTENSIONS:
         return EXTENSIONS[mime]
 
@@ -57,7 +57,7 @@ def _get_ext(mime: str) -> str:
     return extensions[0].strip('.')
 
 
-def _get_mime(ext: str):
+def get_mime(ext: str):
     if ext in MIMES:
         return MIMES[ext]
 
@@ -74,7 +74,7 @@ class Icon:
     @property
     def as_path(self) -> str:
         """Return a filesystem path for the given icon."""
-        ext = _get_ext(self.mime)
+        ext = get_ext(self.mime)
         return str(IMAGE_FOLDER / f'{self.key}_{self.icon_hash}.{ext}')
 
     @property
@@ -83,7 +83,7 @@ class Icon:
 
     @property
     def extension(self) -> str:
-        return _get_ext(self.mime)
+        return get_ext(self.mime)
 
 
 class ImageError(Exception):
@@ -194,6 +194,65 @@ def _try_unlink(path: str):
         pass
 
 
+async def resize_gif(raw_data: bytes, target: tuple) -> tuple:
+    """Resize a GIF image."""
+    # generate a temporary file to call gifsticle to and from.
+    input_fd, input_path = tempfile.mkstemp(suffix='.gif')
+    _, output_path = tempfile.mkstemp(suffix='.gif')
+
+    input_handler = os.fdopen(input_fd, 'wb')
+
+    # make sure its valid image data
+    data_fd = BytesIO(raw_data)
+    image = Image.open(data_fd)
+    image.close()
+
+    log.info('resizing a GIF from {} to {}',
+             image.size, target)
+
+    # insert image info on input_handler
+    # close it to make it ready for consumption by gifsicle
+    input_handler.write(raw_data)
+    input_handler.close()
+
+    # call gifsicle under subprocess
+    log.debug('input: {}', input_path)
+    log.debug('output: {}', output_path)
+
+    process = await asyncio.create_subprocess_shell(
+        f'gifsicle --resize {target[0]}x{target[1]} '
+        f'{input_path} > {output_path}',
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+    )
+
+    # run it, etc.
+    out, err = await process.communicate()
+
+    log.debug('out + err from gifsicle: {}', out + err)
+
+    # write over an empty data_fd
+    data_fd = BytesIO()
+    output_handler = open(output_path, 'rb')
+    data_fd.write(output_handler.read())
+
+    # close unused handlers
+    output_handler.close()
+
+    # delete the files we created with mkstemp
+    _try_unlink(input_path)
+    _try_unlink(output_path)
+
+    # reseek, save to raw_data, reseek again.
+    # TODO: remove raw_data altogether as its inefficient
+    # to have two representations of the same bytes
+    data_fd.seek(0)
+    raw_data = data_fd.read()
+    data_fd.seek(0)
+
+    return data_fd, raw_data
+
+
 class IconManager:
     """Main icon manager."""
     def __init__(self, app):
@@ -201,7 +260,7 @@ class IconManager:
         self.storage = app.storage
 
     async def _convert_ext(self, icon: Icon, target: str):
-        target_mime = _get_mime(target)
+        target_mime = get_mime(target)
         log.info('converting from {} to {}', icon.mime, target_mime)
 
         target_path = IMAGE_FOLDER / f'{icon.key}_{icon.icon_hash}.{target}'
@@ -257,65 +316,6 @@ class IconManager:
         return await self.generic_get(
             'guild', guild_id, icon_hash, **kwargs)
 
-    async def _resize_gif(self, scope: str, key,
-                          raw_data: bytes, target: tuple) -> tuple:
-        """Resize a GIF image."""
-        # generate a temporary file to call gifsticle to and from.
-        input_fd, input_path = tempfile.mkstemp(suffix='.gif')
-        _, output_path = tempfile.mkstemp(suffix='.gif')
-
-        input_handler = os.fdopen(input_fd, 'wb')
-
-        # make sure its valid image data
-        data_fd = BytesIO(raw_data)
-        image = Image.open(data_fd)
-        image.close()
-
-        log.info('resizing a GIF from {} to {}',
-                 image.size, target)
-
-        # insert image info on input_handler
-        # close it to make it ready for consumption by gifsicle
-        input_handler.write(raw_data)
-        input_handler.close()
-
-        # call gifsicle under subprocess
-        log.debug('input: {}', input_path)
-        log.debug('output: {}', output_path)
-
-        process = await asyncio.create_subprocess_shell(
-            f'gifsicle --resize {target[0]}x{target[1]} '
-            f'{input_path} > {output_path}',
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
-        )
-
-        # run it, etc.
-        await process.communicate()
-
-        # write over an empty data_fd
-        data_fd = BytesIO()
-        output_handler = open(output_path, 'rb')
-        data_fd.write(output_handler.read())
-
-        # close unused handlers
-        output_handler.close()
-
-        # delete the files we created with mkstemp
-        _try_unlink(input_path)
-        _try_unlink(output_path)
-
-        # reseek, save to raw_data, reseek again.
-        # TODO: remove raw_data altogether as its inefficient
-        # to have two representations of the same bytes
-        data_fd.seek(0)
-        raw_data = data_fd.read()
-        data_fd.seek(0)
-
-
-        return data_fd, raw_data
-
-
     async def put(self, scope: str, key: str,
                   b64_data: str, **kwargs) -> Icon:
         """Insert an icon."""
@@ -328,7 +328,7 @@ class IconManager:
         data_fd = BytesIO(raw_data)
 
         # get an extension for the given data uri
-        extension = _get_ext(mime)
+        extension = get_ext(mime)
 
         if 'bsize' in kwargs and len(raw_data) > kwargs['bsize']:
             return _invalid(kwargs)
@@ -336,8 +336,7 @@ class IconManager:
         # size management is different for gif files
         # as they're composed of multiple frames.
         if 'size' in kwargs and mime == 'image/gif':
-            data_fd, raw_data = await self._resize_gif(
-                scope, key, raw_data, kwargs['size'])
+            data_fd, raw_data = await resize_gif(raw_data, kwargs['size'])
         elif 'size' in kwargs:
             image = Image.open(data_fd)
 

litecord/schemas.py

@@ -384,7 +384,7 @@ MEMBER_UPDATE = {
 
 
 MESSAGE_CREATE = {
-    'content': {'type': 'string', 'minlength': 1, 'maxlength': 2000},
+    'content': {'type': 'string', 'minlength': 0, 'maxlength': 2000},
     'nonce': {'type': 'snowflake', 'required': False},
     'tts': {'type': 'boolean', 'required': False},
 

litecord/storage.py

@@ -64,8 +64,9 @@ def _filter_recipients(recipients: List[Dict[str, Any]], user_id: int):
 
 class Storage:
     """Class for common SQL statements."""
-    def __init__(self, db):
-        self.db = db
+    def __init__(self, app):
+        self.app = app
+        self.db = app.db
         self.presence = None
 
     async def fetchrow_with_json(self, query: str, *args):
@@ -636,6 +637,52 @@ class Storage:
         # they were defined in the first loop.
         return list(map(react_stats.get, emoji))
 
+    async def get_attachments(self, message_id: int) -> List[Dict[str, Any]]:
+        """Get a list of attachment objects tied to the message."""
+        attachment_ids = await self.db.fetch("""
+        SELECT id
+        FROM attachments
+        WHERE message_id = $1
+        """, message_id)
+
+        attachment_ids = [r['id'] for r in attachment_ids]
+
+        res = []
+
+        for attachment_id in attachment_ids:
+            row = await self.db.fetchrow("""
+            SELECT id::text, message_id, channel_id,
+                   filename, filesize, image, height, width
+            FROM attachments
+            WHERE id = $1
+            """, attachment_id)
+
+            drow = dict(row)
+
+            drow.pop('message_id')
+            drow.pop('channel_id')
+
+            drow['size'] = drow['filesize']
+            drow.pop('size')
+
+            # construct attachment url
+            proto = 'https' if self.app.config['IS_SSL'] else 'http'
+            main_url = self.app.config['MAIN_URL']
+
+            drow['url'] = (f'{proto}://{main_url}/attachments/'
+                           f'{row["channel_id"]}/{row["message_id"]}/'
+                           f'{row["filename"]}')
+
+            # NOTE: since the url comes from the instance itself
+            # i think proxy_url=url is valid.
+            drow['proxy_url'] = drow['url']
+
+            # TODO: url, proxy_url
+
+            res.append(drow)
+
+        return res
+
     async def get_message(self, message_id: int, user_id=None) -> Dict:
         """Get a single message's payload."""
         row = await self.fetchrow_with_json("""
@@ -703,7 +750,7 @@ class Storage:
         res.pop('author_id')
 
         # TODO: res['attachments']
-        res['attachments'] = []
+        res['attachments'] = await self.get_attachments(message_id)
 
         # TODO: res['member'] for partial member data
         #  of the author

manage/cmd/migration/scripts/10_add_attachments_table.sql

@@ -0,0 +1,15 @@
+CREATE TABLE IF NOT EXISTS attachments (
+    id bigint PRIMARY KEY,
+
+    channel_id bigint REFERENCES channels (id),
+    message_id bigint REFERENCES messages (id),
+
+    filename text NOT NULL,
+    filesize integer,
+
+    image boolean DEFAULT FALSE,
+
+    -- only not null if image=true
+    height integer DEFAULT NULL,
+    width integer DEFAULT NULL
+);

run.py

@@ -34,7 +34,8 @@ import config
 
 from litecord.blueprints import (
     gateway, auth, users, guilds, channels, webhooks, science,
-    voice, invites, relationships, dms, icons, nodeinfo, static
+    voice, invites, relationships, dms, icons, nodeinfo, static,
+    attachments
 )
 
 # those blueprints are separated from the "main" ones
@@ -131,6 +132,7 @@ def set_blueprints(app_):
         fake_store: None,
 
         icons: -1,
+        attachments: -1,
         nodeinfo: -1,
         static: -1
     }
@@ -218,7 +220,7 @@ def init_app_managers(app):
     app.ratelimiter = RatelimitManager(app.config.get('_testing'))
     app.state_manager = StateManager()
 
-    app.storage = Storage(app.db)
+    app.storage = Storage(app)
     app.user_storage = UserStorage(app.storage)
 
     app.icons = IconManager(app)

schema.sql

@@ -52,6 +52,26 @@ CREATE TABLE IF NOT EXISTS instance_invites (
 );
 
 
+-- main attachments table
+CREATE TABLE IF NOT EXISTS attachments (
+    id bigint PRIMARY KEY,
+
+    -- keeping channel_id and message_id
+    -- make a way "better" attachment url.
+    channel_id bigint REFERENCES channels (id),
+    message_id bigint REFERENCES messages (id),
+
+    filename text NOT NULL,
+    filesize integer,
+
+    image boolean DEFAULT FALSE,
+
+    -- only not null if image=true
+    height integer DEFAULT NULL,
+    width integer DEFAULT NULL,
+);
+
+
 CREATE TABLE IF NOT EXISTS icons (
     -- can be 'user', 'guild', 'emoji'
     scope text NOT NULL,
@@ -619,12 +639,6 @@ CREATE TABLE IF NOT EXISTS messages (
     message_type int NOT NULL
 );
 
-CREATE TABLE IF NOT EXISTS message_attachments (
-    message_id bigint REFERENCES messages (id),
-    attachment bigint REFERENCES attachments (id),
-    PRIMARY KEY (message_id, attachment)
-);
-
 CREATE TABLE IF NOT EXISTS message_reactions (
     message_id bigint REFERENCES messages (id),
     user_id bigint REFERENCES users (id),