diff --git a/.gitignore b/.gitignore index 062ffb8..e4064a2 100644 --- a/.gitignore +++ b/.gitignore @@ -105,5 +105,6 @@ venv.bak/ config.py images/* +attachments/* .DS_Store diff --git a/attachments/.gitkeep b/attachments/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/litecord/blueprints/__init__.py b/litecord/blueprints/__init__.py index 701bbd2..f43f9ca 100644 --- a/litecord/blueprints/__init__.py +++ b/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'] diff --git a/litecord/blueprints/attachments.py b/litecord/blueprints/attachments.py new file mode 100644 index 0000000..2ece3e0 --- /dev/null +++ b/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 . + +""" + +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' + '///', + 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) diff --git a/litecord/blueprints/channel/messages.py b/litecord/blueprints/channel/messages.py index f1c8a5c..6d30c05 100644 --- a/litecord/blueprints/channel/messages.py +++ b/litecord/blueprints/channel/messages.py @@ -18,11 +18,12 @@ along with this program. If not, see . """ 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('//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: diff --git a/litecord/images.py b/litecord/images.py index 1cb9eb2..d1395aa 100644 --- a/litecord/images.py +++ b/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) diff --git a/litecord/schemas.py b/litecord/schemas.py index c631572..d21ef54 100644 --- a/litecord/schemas.py +++ b/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}, diff --git a/litecord/storage.py b/litecord/storage.py index 6b99dd3..fe4f317 100644 --- a/litecord/storage.py +++ b/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 diff --git a/manage/cmd/migration/scripts/10_add_attachments_table.sql b/manage/cmd/migration/scripts/10_add_attachments_table.sql new file mode 100644 index 0000000..b1f979d --- /dev/null +++ b/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 +); diff --git a/run.py b/run.py index addd7f6..adb01f4 100644 --- a/run.py +++ b/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) diff --git a/schema.sql b/schema.sql index e3df14d..72d68e1 100644 --- a/schema.sql +++ b/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),