Commit 5480a669 authored by Luna's avatar Luna 😻

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
parent e54bcc31
......@@ -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'])
async def update_guild(guild_id):
@bp.route('/<int:guild_id>', 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(
......
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'])
async def get_guild_icon(guild_id: int, icon_hash: str, ext: str):
pass
@bp.route('/icons/<int:guild_id>/<icon_file>', 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/<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):
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/<int:application_id>/<icon_hash>.<ext>')
......
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
......@@ -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 = {
......
-- 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;
......@@ -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,
......
......@@ -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,
-- can be a user snowflake, guild snowflake or
-- emoji snowflake
key text,
-- sha512(file)
hash text NOT NULL,
mimetype text NOT NULL,
-- sha256 of the icon
hash text UNIQUE NOT NULL,
-- path to the file system
fspath text 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)
);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment