globalforest/cncbot.py

import pydle
import asyncio
from asyncio import new_event_loop, gather, get_event_loop, sleep
import random
import string
import time

CNC_HOST='dark.deflax.net'
CNC_CHANNEL='#izba'
ADMIN_NICKNAMES = [ 'afx' ]
MOTD='] Forest CnC ready!'

def get_random_string(length):
    letters = string.ascii_lowercase
    result_str = ''.join(random.choice(letters) for i in range(length))
    return result_str

class ForestBot(pydle.Client):  
    async def on_connect(self):
        await super().on_connect()
        await self.join(CNC_CHANNEL)
        await self.message(CNC_CHANNEL, MOTD)
        
    async def is_admin(self, nickname):
        """
        Check whether or not a user has administrative rights for this bot.
        This is a blocking function: use a coroutine to call it.
        See pydle's documentation on blocking functionality for details.
        """
        admin = False
        # Check the WHOIS info to see if the source has identified with NickServ.
        # This is a blocking operation, so use yield.
        if nickname in ADMIN_NICKNAMES:
            info = await self.whois(nickname)
            print('] Detected privileged request from {}'.format(info))
            admin = info['identified']
        return admin

    async def on_message(self, target, source, message):
        await super().on_message(target, source, message)
        
        # Print help
        if message.startswith('help {}'.format(self.nickname)):
            await self.message(target, '] cnc usage:')
            await self.message(target, '] help - displays help')
            await self.message(target, '] id - cnc user info')
            await self.message(target, '] dump - dump cnc log')

        # Show user info
        if message.startswith('id {}'.format(self.nickname)):
            admin = await self.is_admin(source)
            if admin:
                await self.message(target, '] You are an administrator. :)'.format(source))
            else:
                await self.message(target, '] You are NOT an administrator. :('.format(source))
                
        if message.startswith('dump {}'.format(self.nickname)):
            admin = await self.is_admin(source)
            if admin:
                await self.dumplog(target)
            else:
                await self.message(target, '] You are NOT an administrator. :('.format(source))
                    
    async def dumplog(self, target):
        sleep_time = 1
        while True:
            msg = self.queue.get()
            await asyncio.sleep(sleep_time)
            await self.message(target, msg)
            if len(msg) == 0:
                break
    
    def run(self, *args, **kwargs):
        self.queue = kwargs['extqueue']
        
        loop = self.eventloop
        #loop.run_until_complete(self.connect(*args, **kwargs))
        loop.run_until_complete(self.connect(*args))
        try:
            loop.run_forever()
        finally:
            loop.stop()
            
def CnCApp(queue):
    botnick = 'cnc_' + get_random_string(8)
    botnick_fallback = [botnick + '_re']
    cnc = ForestBot(botnick, botnick_fallback, botnick, botnick, None)
    cnc.run(CNC_HOST, tls=False, tls_verify=False, extqueue=queue)

if __name__ == '__main__':
    CnCApp()