105 lines
4.5 KiB
Python
105 lines
4.5 KiB
Python
import asyncio,websockets,logging,ujson,argparse
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
# Command line specific arguments
|
|
parser.add_argument('--token', help='Bot Authentication Token', default=None)
|
|
parser.add_argument('--sequence', help='Bot Authentication Sequence', default=None)
|
|
parser.add_argument('--session', help='Bot Authentication Session', default=None)
|
|
|
|
class tinyclient():
|
|
def __init__ (self, *args, **kwargs):
|
|
self.opcode_map = {0:'dispatch',1:'heartbeat',2:'identify',3:'presence',4:'voice_state',5:'voice_ping',6:'resume',7:'reconnect',8:'request_members',9:'invalidate_session',10:'hello',11:'heartbeat_ack',12:'guild_sync'}
|
|
self.last_sequence = kwargs.pop('sequence',None)
|
|
self.session_id = kwargs.pop('session',None)
|
|
self.token=kwargs.pop('token',None)
|
|
self.user=None
|
|
self.scheduler = AsyncIOScheduler(timezone='utc')
|
|
self.scheduler.start()
|
|
self.websocket = None
|
|
self.log = logging.getLogger(__name__)
|
|
#self.log.setLevel(logging.DEBUG)
|
|
#self.log.addHandler(logging.StreamHandler())
|
|
self.loop = asyncio.get_event_loop()
|
|
def event(self, coro):
|
|
if not asyncio.iscoroutinefunction(coro):
|
|
raise Exception('event registered must be a coroutine function')
|
|
setattr(self, coro.__name__, coro)
|
|
self.log.debug('{0.__name__} has successfully been registered as an event'.format(coro))
|
|
return coro
|
|
def dispatch(self, event, *args, **kwargs):
|
|
self.log.debug('Dispatching event {}'.format(event))
|
|
method = 'on_' + event
|
|
handler = 'handle_' + event
|
|
|
|
if hasattr(self, handler):
|
|
o = getattr(self, handler)
|
|
asyncio.async(o(*args, **kwargs), loop=self.loop)
|
|
|
|
if hasattr(self, method):
|
|
o = getattr(self, method)
|
|
asyncio.async(o(*args, **kwargs), loop=self.loop)
|
|
async def next_message(self):
|
|
message = await self.websocket.recv()
|
|
o = ujson.loads(message)
|
|
self.last_sequence = o.get('s')
|
|
return o
|
|
async def run(self,*args,**kwargs):
|
|
async with websockets.connect('wss://gateway.discord.gg/?encoding=json&v=6') as self.websocket:
|
|
while True:
|
|
o = await self.next_message()
|
|
event_name = self.opcode_map[o.get('op')]
|
|
self.dispatch(event_name, o)
|
|
|
|
async def handle_hello(self,data):
|
|
interval = int(data.get('d').get('heartbeat_interval')) / 1000
|
|
self.scheduler.add_job(self.websocket.send, 'interval', seconds=interval, args=(ujson.dumps({"op": 1, "d": self.last_sequence}),))
|
|
|
|
if self.last_sequence and self.session_id:
|
|
resume = {"op":6,"d":{"token":self.token, "session_id":self.session_id, "seq":self.last_sequence}}
|
|
await self.websocket.send(ujson.dumps(resume))
|
|
else:
|
|
identify = {"op":2,"d":{"token":self.token,"properties":{"os":"Linux"},"large_threshold":250,"synced_guilds":[],"presence":{"status":"online","since":0,"afk":False,"game":None},"compress":False}}
|
|
await self.websocket.send(ujson.dumps(identify))
|
|
async def handle_dispatch(self, data):
|
|
if data.get('t') == 'READY':
|
|
self.session_id = data.get('d').get('session_id')
|
|
self.user = data.get('d').get('user')
|
|
|
|
if data.get('t'):
|
|
event_name = data.get('t').lower()
|
|
self.dispatch(event_name, data)
|
|
def bot_main(run=False):
|
|
|
|
# Parse out all our command line arguments
|
|
args = parser.parse_args()
|
|
logging.info("token=%s, sequence=%s, session=%s,",args.token,args.sequence,args.session)
|
|
|
|
|
|
# create a new bot-class instance
|
|
t = tinyclient(token=args.token,sequence=args.sequence,session=args.session)
|
|
|
|
@t.event
|
|
async def on_message_delete(*args):
|
|
data, = args
|
|
id = data.get('d').get('id')
|
|
|
|
@t.event
|
|
async def on_message_update(*args):
|
|
data, = args
|
|
id = data.get('d').get('id')
|
|
|
|
@t.event
|
|
async def on_message_create(*args):
|
|
data, = args
|
|
|
|
asyncio.get_event_loop().run_until_complete(t.run())
|
|
asyncio.get_event_loop().run_forever()
|
|
|
|
if __name__ == "__main__":
|
|
# configure basic logging (also muting noisy libs)
|
|
logging.basicConfig(level=logging.INFO,format='%(levelname)-6s%(name)-8s%(asctime)s - %(message)s')
|
|
[logging.getLogger(x).setLevel(logging.ERROR) for x in ['apscheduler']]
|
|
|
|
bot_main(True) |