server.py

import os
import re
import time
from glob import glob
import traceback
import hmac
import asyncio
import mimetypes
from aiohttp import web, ClientSession, BasicAuth
from db import Database, USERS_API
from responses import *

DEFAULT_PORT = 8888
COMMENTS_API = 'https://scratch.mit.edu/site-api/comments/user/{}/\
?page=1&salt={}'
USERNAME_REGEX = re.compile('^[A-Za-z0-9_-]{3,20}$')
COMMENTS_REGEX = re.compile(r"""<div id="comments-\d+" class="comment +" data-comment-id="\d+">.*?<div class="actions-wrap">.*?<div class="name">\s+<a href="/users/([_a-zA-Z0-9-]+)">\1</a>\s+</div>\s+<div class="content">(.*?)</div>""", re.S)

class Server:
    def __init__(self, hook_secret, discord_hook, name=None):
        self.app = web.Application(middlewares=[self.errors])
        self.app.add_routes([
            web.put('/verify/{username}', self.verify),
            web.post('/verify/{username}', self.verified),
            web.delete('/verify/{username}', self.unverify),
            web.post('/users/{username}/login', self.login),
            web.post('/users/{username}/finish-login', self.finish_login),
            web.post('/users/{username}/logout', self.logout_user),
            web.get('/session/{session}', self.get_user),
            web.put('/session/{session}', self.put_user),
            web.patch('/session/{session}', self.reset_token),
            web.delete('/session/{session}', self.del_user),
            web.post('/session/{session}/logout', self.logout),
            web.get('/usage', self.logs),
            web.get('/usage/{logid}', self.log),
            web.post('/webhook', self.gh_hook),
            web.get('/site/{path:.*}', self.file_handler),
            web.get('/site/', self.file_handler),
            web.get('/site', self.file_handler),
            web.get('/docs/{path:.*}', self.docs_handler),
            web.get('/docs/', self.docs_handler),
            web.get('/docs', self.docs_handler),
            web.post('/debug/{dibool}', self.set_debug),
            web.view('/{path:.*}', self.not_found)
        ])
        self.session = ClientSession()
        self.db = Database(self.session)
        self.hook_secret = hook_secret.encode()
        self.discord_hook = discord_hook
        self.name = name
        self.debug = False
        self._debug = False
        self.debug_pass = False

    @web.middleware
    async def errors(self, request, handler):
        print('%s %s %s' % (
            time.strftime('%Y-%m-%dT%H:%M:%SZ'),
            request.method,
            request.path_qs
        ))
        try:
            return await handler(request)
        except web.HTTPException as exc:
            raise
        except Exception as exc:
            if self._debug:
                print(traceback.format_exc())
                raise web.HTTPServerError()
            await self.session.post(self.discord_hook, json={
                'username': '{}ScratchVerifier Errors'.format(
                    (self.name + "'s ") if self.name else ''
                ),
                'embeds': [{
                    'color': 0xff0000,
                    'title': '500 Response Sent',
                    'fields': [
                        {'name': 'Request Path',
                         'value': '`%s %s`' % (request.method, request.path_qs),
                         'inline': False},
                        {'name': 'Error traceback',
                         'value': '```%s```' % traceback.format_exc(),
                         'inline': False}
                    ]
                }]
            })
            raise

    async def run(self, port=DEFAULT_PORT, debug=False):
        self.runner = web.AppRunner(self.app)
        self._debug = debug
        await self.runner.setup()
        site = web.TCPSite(self.runner, '0.0.0.0', port)
        await site.start()

    async def stop(self):
        await self.runner.cleanup()
        await self.session.close()
        await self.db.close()

    async def _wakeup(self):
        files = glob(os.path.join(os.path.dirname(__file__), '*.py'))
        files.append(os.path.join(os.path.dirname(__file__),
                                  'sql', 'startup.sql'))
        files = {f: os.path.getmtime(f) for f in files}
        while 1:
            try:
                for i in files:
                    if os.path.getmtime(i) > files[i]:
                        return
                await asyncio.sleep(1)
            except:
                return

    def run_sync(self, port=DEFAULT_PORT, debug=False):
        loop = asyncio.get_event_loop()
        try:
            loop.run_until_complete(self.run(port, debug))
            loop.run_until_complete(self._wakeup())
        except KeyboardInterrupt:
            pass
        finally:
            loop.run_until_complete(self.stop())

    async def check_token(self, request):
        if self.debug:
            return 0
        if 'Authorization' not in request.headers:
            raise web.HTTPUnauthorized()
        auth = BasicAuth.decode(request.headers['Authorization'])
        client_id = int(auth.login)
        if not await self.db.client_matches(client_id, auth.password):
            raise web.HTTPUnauthorized()
        return client_id

    async def check_username(self, request):
        username = request.match_info.get('username', '')
        if not re.match(USERNAME_REGEX, username):
            raise web.HTTPBadRequest()
        if self.debug:
            return username.casefold()
        async with self.session.get(USERS_API.format(username)) as resp:
            if resp.status != 200:
                raise web.HTTPNotFound()
        return username.casefold()

    async def verify(self, request):
        client_id = await self.check_token(request)
        username = await self.check_username(request)
        code = await self.db.start_verification(client_id, username)
        return web.json_response(Verification(
            code=code, username=username
        ).d())

    async def _verified(self, client_id, username):
        if self.debug_pass:
            await self.db.end_verification(client_id, username, True)
            return True
        code = await self.db.get_code(client_id, username)
        if not code:
            raise web.HTTPNotFound()
        if self.debug:
            await self.db.end_verification(client_id, username, False)
            return False
        async with self.session.get(COMMENTS_API.format(
            username, int(time.time())
        )) as resp:
            if resp.status != 200:
                raise web.HTTPNotFound() #likely banned or something
            data = await resp.text()
        data = data.strip()
        if not data:
            await self.db.end_verification(client_id, username, False)
            return False #no comments at all, failed
        for m in re.finditer(COMMENTS_REGEX, data):
            if m.group(1).casefold() != username:
                continue
            if m.group(2).strip() == code:
                break
        else:
            await self.db.end_verification(client_id, username, False)
            return False #nothing was the code, failed
        await self.db.end_verification(client_id, username, True)
        return True

    async def verified(self, request):
        client_id = await self.check_token(request)
        username = await self.check_username(request)
        if await self._verified(client_id, username):
            raise web.HTTPNoContent()
        else:
            raise web.HTTPForbidden()

    async def unverify(self, request):
        client_id = await self.check_token(request)
        username = await self.check_username(request)
        await self.db.end_verification(client_id, username, -1)
        raise web.HTTPNoContent()

    async def login(self, request):
        username = await self.check_username(request)
        code = await self.db.start_verification(0, username)
        return web.json_response(Verification(
            code=code, username=username
        ).d())

    async def finish_login(self, request):
        username = await self.check_username(request)
        if await self._verified(0, username):
            return web.json_response(Session(
                session=await self.db.new_session(username)
            ).d())
        else:
            raise web.HTTPUnauthorized()

    async def logout_user(self, request):
        session_id = await self.check_session(request)
        data = User(**(await self.db.get_client(session_id)))
        username = await self.check_username(request)
        if username != data.username:
            raise web.HTTPForbidden()
        await self.db.logout_user(username)

    async def check_session(self, request):
        if self.debug:
            return int(request.match_info['session'])
        session = request.match_info.get('session',
                                         request.query.get('session', ''))
        if not (session or session.strip()):
            raise web.HTTPUnauthorized()
        try:
            session_id = int(session)
        except ValueError:
            raise web.HTTPUnauthorized() from None
        if await self.db.get_expired(session_id):
            raise web.HTTPUnauthorized()
        return session_id

    async def get_user(self, request):
        session_id = await self.check_session(request)
        data = await self.db.get_client(session_id)
        if data is None:
            raise web.HTTPNotFound()
        return web.json_response(data)

    async def put_user(self, request):
        session_id = await self.check_session(request)
        data = await self.db.get_client(session_id)
        if data is not None:
            raise web.HTTPConflict()
        if self.debug:
            session_id = 0 #prevent fake non-conflict debug IDs from registering
        data = await self.db.new_client(session_id)
        return web.json_response(data)

    async def reset_token(self, request):
        session_id = await self.check_session(request)
        await self.db.reset_token(session_id)
        return await self.get_user(request)

    async def del_user(self, request):
        session_id = await self.check_session(request)
        await self.db.del_client(session_id)
        raise web.HTTPNoContent()

    async def logout(self, request):
        session_id = await self.check_session(request)
        await self.db.logout(session_id)
        raise web.HTTPNoContent()

    async def logs(self, request):
        params = {}
        for k in {'limit', 'start', 'end', 'before',
                  'after', 'client_id', 'type'}:
            if k in request.query:
                try:
                    params[k] = int(request.query[k])
                except ValueError:
                    raise web.HTTPBadRequest() from None
        if 'username' in request.query:
            params['username'] = request.query['username']
        if params.get('limit', 100) > 500:
            raise web.HTTPForbidden()
        return web.json_response(await self.db.get_logs(**params))

    async def log(self, request):
        log_id = request.match_info.get('logid', None)
        try:
            log_id = int(log_id)
        except ValueError:
            raise web.HTTPNotFound() from None
        data = await self.db.get_log(log_id)
        if data is None:
            raise web.HTTPNotFound()
        return web.json_response(data)

    async def gh_hook(self, request):
        if 'X-Hub-Signature' not in request.headers:
            raise web.HTTPUnauthorized()
        algo, hash = request.headers['X-Hub-Signature'].split('=')
        digest = hmac.new(self.hook_secret, await request.read(), algo)
        digest = digest.hexdigest()
        if not hmac.compare_digest(digest, hash):
            raise web.HTTPUnauthorized()
        data = await request.json()
        if data['ref'] != 'refs/heads/master':
            raise web.HTTPNoContent() #success, but no action
        os.system('cd {} && git pull --rebase origin'.format(
            os.path.dirname(__file__)
        ))
        raise web.HTTPNoContent()

    async def set_debug(self, request):
        if not self._debug:
            raise web.HTTPNotFound()
        val = int(request.match_info['dibool'])
        self.debug, self.debug_pass = val & 1, val & 2
        raise web.HTTPNoContent()

    async def _file_handler(self, request, WEB_ROOT):
        WEB_ROOT = os.path.join(os.path.dirname(
            os.path.dirname(__file__)
        ), WEB_ROOT)
        PATH = request.match_info.get('path', 'index.html') or 'index.html'
        if '.' not in PATH.split('/')[-1]:
            PATH += '/index.html'
        FILE = os.path.join(WEB_ROOT, PATH)
        if request.if_modified_since:
            if os.path.getmtime(FILE) <= \
                    request.if_modified_since.timestamp():
                raise web.HTTPNotModified()
        if request.if_unmodified_since:
            if os.path.getmtime(FILE) > \
                    request.if_unmodified_since.timestamp():
                raise web.HTTPPreconditionFailed()
        try:
            with open(FILE, 'rb') as f:
                range = request.http_range
                if range.stop:
                    if range.start:
                        f.seek(range.start)
                        rsize = range.stop - range.start
                    else:
                        rsize = range.stop
                elif range.start:
                    f.seek(range.start)
                    rsize = -1
                else:
                    rsize = -1
                data = f.read(rsize)
            ct = mimetypes.guess_type(PATH)
            return web.Response(body=data, content_type=ct[0], charset=ct[1])
        except FileNotFoundError:
            raise web.HTTPNotFound()

    async def file_handler(self, request):
        return await self._file_handler(request, 'public')

    async def docs_handler(self, request):
        return await self._file_handler(request, 'docs')

    async def not_found(self, request):
        raise web.HTTPNotFound()