index.py

import asyncio
import websockets
import json
import os
import signal
import logging
import time

logfmt = '%(asctime)s - %(levelname)s - %(message)s'
loghdlrs = [logging.FileHandler('index.log'), logging.StreamHandler()]
logging.basicConfig(level=logging.INFO, format=logfmt, handlers=loghdlrs)
logging.basicConfig(level=logging.ERROR, format=logfmt, handlers=loghdlrs)
logging.basicConfig(level=logging.CRITICAL, format=logfmt, handlers=loghdlrs)

if os.getenv('IDEBUG') == True:
    logging.basicConfig(level=logging.DEBUG)

class WSS:
    def __init__(self):
        self.STATE = 0
        self.WSS_PORT = 2112

        self.UID_SOCK = dict()
        self.ADDR_UID = dict()

    async def start(self, stop):
        async with websockets.serve(self.ws, '0.0.0.0', self.WSS_PORT):
            logging.info(f'Starting server on {self.WSS_PORT} port')
            await stop

    async def checkuser(self, user, asker, sock):
        if user in self.ADDR_UID:
            msg = {
                'm': 'get',
                'q': 'you',
                'u': asker
            }

            jmsg = json.dumps(msg)
            logging.info(f'> {jmsg}')
            await self.UID_SOCK[self.ADDR_UID[user]].send(jmsg)
        else:
            res = {
                'm': 'info',
                'q': 'checkuser',
                'u': user,
                'c': 404,
            }
            jres = json.dumps(res)
            logging.info(f'> {jres}')
            await sock.send(jres)

    async def senduser(self, asker, name, status):

        if asker in self.UID_SOCK: 

            res = {
                'm': 'info',
                'q': 'checkuser',
                'c': 200,
                'res': {
                    'name': name,
                    'status': status
                }
            }
            jres = json.dumps(res)
            logging.info(f'> {jres}')
            await self.UID_SOCK[asker].send(jres)

    async def sendmsg(self, addr, sender, ts, pack):
        if addr in self.ADDR_UID:

            msg = {
                'm': 'put',
                'q': 'msg',
                'user': sender,
                'pack': pack
            }
            jmsg = json.dumps(msg)
            logging.info(f'> {jmsg}')
            await self.UID_SOCK[self.ADDR_UID[addr]].send(jmsg)

            c = 200
            #await asyncio.wait([self.USER_SOCK[_user].send(jmsg) for _user in self.PUB_USER[pub]])
        else:
            c = 404

        res = {
            'm': 'info',
            'q': 'sendmsg',
            'd': ts,
            'c': c,
            'pack': pack
        }
        return json.dumps(res)

    async def checkin(self, uid, sock, addr):

        if addr in self.ADDR_UID:
            c = 409
        else:
            self.ADDR_UID.update({addr: uid})
            c = 201

        res = f'{{"m": "info", "q": "checkin", "c": {c}}}'

        logging.info(f'CHECKIN {addr} {uid}')

        return res

    async def checkout(self, addr):

        if addr in self.ADDR_UID:
            self.ADDR_UID.pop(addr, None)
            logging.info(f'CHECKOUT {addr}')
        else:
            pass     

    async def ws(self, SOCK, path=None):

        UID = int(time.time() * 1000)
        self.UID_SOCK.update({UID: SOCK})
        logging.info(f'OPEN CONNECTION FROM {UID}')

        ADDR = str()

        res = {
            'm': "info", 
            'q': "srv", 
            'c': 100, 
            'res' : {
                'pub': "123123123",
                'uid': UID,
                'online': len(self.UID_SOCK)
            }
        }

        await SOCK.send(json.dumps(res))

        try:
            async for sock in SOCK:
                #print(websocket.closed)
                #print(websocket.open)
                req = json.loads(sock)
                logging.info(f'< {req}')
                method = req['m']
                query = req['q']
                #print (locals())

                if method == 'put':

                    if query == 'user':
                        ADDR = req['pack']['addr']
                        res = await self.checkin(UID, SOCK, ADDR)
                        logging.info(f'> {res}')
                        await SOCK.send(res)

                    if query == 'msg':

                        if 'n' not in req:
                            req['n'] = 'Unknown'

                        if 'd' not in req:
                            req['d'] = int(time.time() * 1000)

                        if 'u' in req and 'pack' in req:
                            res = await self.sendmsg(req['u'], req['n'], req['d'], req['pack'])

                        else:
                            res = {
                                'm': 'info',
                                'q': 'sendmsg',
                                'd': req['d'],
                                'c': 400
                            }
                            
                        logging.info(f'> {res}')
                        await SOCK.send(res)        

                if method == 'get':

                    if query == 'user':
                        res = await self.checkuser(req['u'], UID, SOCK)

                if method == 'post':

                    if query == 'user':

                        if 'name' not in req:
                            req['name'] = 'Unknown'

                        if 'status' not in req:
                            req['status']= '~'

                        res = await self.senduser(req['u'], req['name'], req['status'])

        except Exception as e:
            logging.error(f'{type(e).__name__}: {e}')
            #cf01b7a83315db38b082e83cf7dbc1fb7fa73d9606f93211a13efba1ed510e07
        
        finally:
            self.UID_SOCK.pop(UID, None)
            await self.checkout(ADDR)
            logging.info(f'CLOSE CONNECTION FROM {UID}')
if __name__ == '__main__':
    try:
        ws = WSS()
        loop = asyncio.get_event_loop()

        stop = asyncio.Future()
        loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)

        server = loop.run_until_complete(ws.start(stop))
    except (KeyboardInterrupt, SystemExit):
        pass