skeleton

- bot skeleton for irc
git clone git://git.acid.vegas/skeleton.git
skeleton.py (10051B)
      1 #!/usr/bin/env python
      2 # irc bot skeleton - developed by acidvegas in python (https://git.acid.vegas/skeleton)
      3 
      4 import argparse
      5 import asyncio
      6 import logging
      7 import logging.handlers
      8 import ssl
      9 import time
     10 
     11 # Settings
     12 cmd_flood = 3 # Delay between bot command usage in seconds (In this case, anything prefixed with a ! is a command)
     13 
     14 # Formatting Control Characters / Color Codes
     15 bold        = '\x02'
     16 italic      = '\x1D'
     17 underline   = '\x1F'
     18 reverse     = '\x16'
     19 reset       = '\x0f'
     20 white       = '00'
     21 black       = '01'
     22 blue        = '02'
     23 green       = '03'
     24 red         = '04'
     25 brown       = '05'
     26 purple      = '06'
     27 orange      = '07'
     28 yellow      = '08'
     29 light_green = '09'
     30 cyan        = '10'
     31 light_cyan  = '11'
     32 light_blue  = '12'
     33 pink        = '13'
     34 grey        = '14'
     35 light_grey  = '15'
     36 
     37 
     38 def color(msg: str, foreground: str, background: str = None) -> str:
     39 	'''
     40 	Color a string with the specified foreground and background colors.
     41 
     42 	:param msg: The string to color.
     43 	:param foreground: The foreground color to use.
     44 	:param background: The background color to use.
     45 	'''
     46 	return f'\x03{foreground},{background}{msg}{reset}' if background else f'\x03{foreground}{msg}{reset}'
     47 
     48 
     49 def ssl_ctx(verify: bool = False, cert_path: str = None, cert_pass: str = None) -> ssl.SSLContext:
     50 	'''
     51 	Create a SSL context for the connection.
     52 
     53 	:param verify: Verify the SSL certificate.
     54 	:param cert_path: The path to the SSL certificate.
     55 	:param cert_pass: The password for the SSL certificate.
     56 	'''
     57 	ctx = ssl.create_default_context() if verify else ssl._create_unverified_context()
     58 	if cert_path:
     59 		ctx.load_cert_chain(cert_path) if not cert_pass else ctx.load_cert_chain(cert_path, cert_pass)
     60 	return ctx
     61 
     62 
     63 class Bot():
     64 	def __init__(self):
     65 		self.nickname = 'skeleton'
     66 		self.username = 'skelly'
     67 		self.realname = 'Developement Bot'
     68 		self.reader   = None
     69 		self.writer   = None
     70 		self.last     = time.time()
     71 
     72 
     73 	async def action(self, chan: str, msg: str):
     74 		'''
     75 		Send an ACTION to the IRC server.
     76 
     77 		:param chan: The channel to send the ACTION to.
     78 		:param msg: The message to send to the channel.
     79 		'''
     80 		await self.sendmsg(chan, f'\x01ACTION {msg}\x01')
     81 
     82 
     83 	async def raw(self, data: str):
     84 		'''
     85 		Send raw data to the IRC server.
     86 
     87 		:param data: The raw data to send to the IRC server. (512 bytes max including crlf)
     88 		'''
     89 		self.writer.write(data[:510].encode('utf-8') + b'\r\n')
     90 
     91 
     92 	async def sendmsg(self, target: str, msg: str):
     93 		'''
     94 		Send a PRIVMSG to the IRC server.
     95 
     96 		:param target: The target to send the PRIVMSG to. (channel or user)
     97 		:param msg: The message to send to the target.
     98 		'''
     99 		await self.raw(f'PRIVMSG {target} :{msg}')
    100 
    101 
    102 	async def connect(self):
    103 		'''Connect to the IRC server.'''
    104 		while True:
    105 			try:
    106 				options = {
    107 					'host'       : args.server,
    108 					'port'       : args.port if args.port else 6697 if args.ssl else 6667,
    109 					'limit'      : 1024, # Buffer size in bytes (don't change this unless you know what you're doing)
    110 					'ssl'        : ssl_ctx() if args.ssl else None,
    111 					'family'     : 10 if args.v6 else 2, # 10 = AF_INET6 (IPv6), 2 = AF_INET (IPv4)
    112 					'local_addr' : args.vhost if args.vhost else None # Can we just leave this as args.vhost?
    113 				}
    114 				self.reader, self.writer = await asyncio.wait_for(asyncio.open_connection(**options), 15) # 15 second timeout
    115 				if args.password:
    116 					await self.raw('PASS ' + args.password) # Rarely used, but IRCds may require this
    117 				await self.raw(f'USER {self.username} 0 * :{self.realname}') # These lines must be sent upon connection
    118 				await self.raw('NICK ' + self.nickname)                      # They are to identify the bot to the server
    119 				while not self.reader.at_eof():
    120 					data = await asyncio.wait_for(self.reader.readuntil(b'\r\n'), 300) # 5 minute ping timeout
    121 					await self.handle(data.decode('utf-8').strip()) # Handle the data received from the IRC server
    122 			except Exception as ex:
    123 				logging.error(f'failed to connect to {args.server} ({str(ex)})')
    124 			finally:
    125 				await asyncio.sleep(30) # Wait 30 seconds before reconnecting
    126 
    127 
    128 	async def eventPRIVMSG(self, data: str):
    129 		'''
    130 		Handle the PRIVMSG event.
    131 
    132 		:param data: The data received from the IRC server.
    133 		'''
    134 		parts  = data.split()
    135 		ident  = parts[0][1:] # nick!user@host
    136 		nick   = parts[0].split('!')[0][1:] # Nickname of the user who sent the message
    137 		target = parts[2] # Channel or user (us) the message was sent to
    138 		msg	   = ' '.join(parts[3:])[1:]
    139 		if target == self.nickname: # Handle private messages
    140 			if ident == 'acidvegas!stillfree@big.dick.acid.vegas': # Admin only command based on ident
    141 				if msg.startswith('!raw') and len(msg.split()) > 1: # Only allow !raw if there is some data
    142 					option = ' '.join(msg.split()[1:]) # Everything after !raw is stored here
    143 					await self.raw(option) # Send raw data to the server FROM the bot
    144 			else:
    145 				await self.sendmsg(nick, 'Do NOT message me!') # Let's ignore anyone PM'ing the bot that isn't the admin
    146 		if target.startswith('#'): # Handle channel messages
    147 			if msg.startswith('!'):
    148 				if time.time() - self.last < cmd_flood: # Prevent command flooding
    149 					if not self.slow: # The self.slow variable is used so that a warning is only issued one time
    150 						self.slow = True
    151 						await self.sendmsg(target, color('Slow down nerd!', red))
    152 				else: # Once we confirm the user isn't command flooding, we can handle the commands
    153 					self.slow = False
    154 					if msg == '!help':
    155 						await self.action(target, 'explodes')
    156 					elif msg == '!ping':
    157 						await self.sendmsg(target, 'Pong!')
    158 					elif msg.startswith('!say') and len(msg.split()) > 1: # Only allow !say if there is something to say
    159 						option = ' '.join(msg.split()[1:]) # Everything after !say is stored here
    160 						await self.sendmsg(target, option)
    161 					self.last = time.time() # Update the last command time if it starts with ! character to prevent command flooding
    162 
    163 
    164 	async def handle(self, data: str):
    165 		'''
    166 		Handle the data received from the IRC server.
    167 
    168 		:param data: The data received from the IRC server.
    169 		'''
    170 		logging.info(data)
    171 		try:
    172 			parts = data.split()
    173 			if data.startswith('ERROR :Closing Link:'):
    174 				raise Exception('BANNED')
    175 			if parts[0] == 'PING':
    176 				await self.raw('PONG ' + parts[1]) # Respond to the server's PING request with a PONG to prevent ping timeout
    177 			elif parts[1] == '001': # RPL_WELCOME
    178 				await self.raw(f'MODE {self.nickname} +B') # Set user mode +B (Bot)
    179 				await self.sendmsg('NickServ', f'IDENTIFY {self.nickname} simps0nsfan420') # Identify to NickServ
    180 				await self.raw('OPER MrSysadmin fartsimps0n1337') # Oper up
    181 				await asyncio.sleep(10) # Wait 10 seconds before joining the channel (required by some IRCds to wait before JOIN)
    182 				if parts.key:
    183 					await self.raw(f'JOIN {args.channel} {args.key}') # Join the channel with the key
    184 				else:
    185 					await self.raw(f'JOIN {args.channel}')
    186 			elif parts[1] == '433': # ERR_NICKNAMEINUSE
    187 				self.nickname += '_' # If the nickname is already in use, append an underscore to the end of it
    188 				await self.raw('NICK ' + self.nickname) # Send the new nickname to the server
    189 			elif parts[1] == 'INVITE':
    190 				target = parts[2]
    191 				chan = parts[3][1:]
    192 				if target == self.nickname: # If we were invited to a channel, join it
    193 					await self.raw(f'JOIN {chan}')
    194 			elif parts[1] == 'KICK':
    195 				chan   = parts[2]
    196 				kicked = parts[3]
    197 				if kicked == self.nickname: # If we were kicked from the channel, rejoin it after 3 seconds
    198 					await asyncio.sleep(3)
    199 					await self.raw(f'JOIN {chan}')
    200 			elif parts[1] == 'PRIVMSG':
    201 				await self.eventPRIVMSG(data) # We put this in a separate function since it will likely be the most used/handled event
    202 		except (UnicodeDecodeError, UnicodeEncodeError):
    203 			pass # Some IRCds allow invalid UTF-8 characters, this is a very important exception to catch
    204 		except Exception as ex:
    205 			logging.exception(f'Unknown error has occured! ({ex})')
    206 
    207 
    208 def setup_logger(log_filename: str, to_file: bool = False):
    209 	'''
    210 	Set up logging to console & optionally to file.
    211 
    212 	:param log_filename: The filename of the log file
    213 	:param to_file: Whether or not to log to a file
    214 	'''
    215 	sh = logging.StreamHandler()
    216 	sh.setFormatter(logging.Formatter('%(asctime)s | %(levelname)9s | %(message)s', '%I:%M %p'))
    217 	if to_file:
    218 		fh = logging.handlers.RotatingFileHandler(log_filename+'.log', maxBytes=250000, backupCount=3, encoding='utf-8') # Max size of 250KB, 3 backups
    219 		fh.setFormatter(logging.Formatter('%(asctime)s | %(levelname)9s | %(filename)s.%(funcName)s.%(lineno)d | %(message)s', '%Y-%m-%d %I:%M %p')) # We can be more verbose in the log file
    220 		logging.basicConfig(level=logging.NOTSET, handlers=(sh,fh))
    221 	else:
    222 		logging.basicConfig(level=logging.NOTSET, handlers=(sh,))
    223 
    224 
    225 
    226 if __name__ == '__main__':
    227 	parser = argparse.ArgumentParser(description="Connect to an IRC server.") # The arguments without -- are required arguments.
    228 	parser.add_argument("server", help="The IRC server address.")
    229 	parser.add_argument("channel", help="The IRC channel to join.")
    230 	parser.add_argument("--password", help="The password for the IRC server.")
    231 	parser.add_argument("--port", type=int, help="The port number for the IRC server.") # Port is optional, will default to 6667/6697 depending on SSL.
    232 	parser.add_argument("--ssl", action="store_true", help="Use SSL for the connection.")
    233 	parser.add_argument("--v4", action="store_true", help="Use IPv4 for the connection.")
    234 	parser.add_argument("--v6", action="store_true", help="Use IPv6 for the connection.")
    235 	parser.add_argument("--key",  default="", help="The key (password) for the IRC channel, if required.")
    236 	parser.add_argument("--vhost", help="The VHOST to use for connection.")
    237 	args = parser.parse_args()
    238 
    239 	print(f"Connecting to {args.server}:{args.port} (SSL: {args.ssl}) and joining {args.channel} (Key: {args.key or 'None'})")
    240 
    241 	setup_logger('skeleton', to_file=True) # Optionally, you can log to a file, change to_file to False to disable this.
    242 
    243 	bot = Bot() # We define this here as an object so we can call it from an outside function if we need to.
    244 
    245 	asyncio.run(bot.connect())