From 234295ff23df2a1a96a87821f51d612029985f45 Mon Sep 17 00:00:00 2001
From: Ben Kurtovic <ben.kurtovic@verizon.net>
Date: Sat, 5 Nov 2011 16:04:16 -0400
Subject: [PATCH] Colors and some other stuff for logging.

---
 bot/classes/connection.py | 10 ++++++---
 bot/config.py             | 53 ++++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/bot/classes/connection.py b/bot/classes/connection.py
index 0b6ca1f..ab97cb5 100644
--- a/bot/classes/connection.py
+++ b/bot/classes/connection.py
@@ -10,7 +10,7 @@ class BrokenSocketException(Exception):
 
 class Connection(object):
     """A class to interface with IRC."""
-    
+
     def __init__(self, host=None, port=None, nick=None, ident=None,
                  realname=None, logger=None):
         self.host = host
@@ -26,7 +26,11 @@ class Connection(object):
     def connect(self):
         """Connect to our IRC server."""
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.sock.connect((self.host, self.port))
+        try:
+            self.sock.connect((self.host, self.port))
+        except socket.error:
+            self.logger.critical("Couldn't connect to IRC server", exc_info=1)
+            exit(1)
         self.send("NICK %s" % self.nick)
         self.send("USER %s %s * :%s" % (self.ident, self.host, self.realname))
 
@@ -68,7 +72,7 @@ class Connection(object):
         message = "".join((chr(1), "ACTION ", msg, chr(1)))
         self.say(target, message)
 
-    def notice(self, target, msg): 
+    def notice(self, target, msg):
         """Send a notice to a target on the server."""
         message = "".join(("NOTICE ", target, " :", msg))
         self.send(message)
diff --git a/bot/config.py b/bot/config.py
index 8c78c0b..02626a9 100644
--- a/bot/config.py
+++ b/bot/config.py
@@ -27,7 +27,7 @@ Additionally, there are functions used in config loading:
 import json
 import logging
 import logging.handlers
-from os import path
+from os import mkdir, path
 
 import blowfish
 
@@ -59,26 +59,36 @@ def _setup_logging():
 
     if metadata.get("enableLogging"):
         hand = logging.handlers.TimedRotatingFileHandler
-        fmt = "[%(asctime)s %(levelname)-8s] %(name)s: %(message)s"
-        formatter = logging.Formatter(fmt=fmt, datefmt="%Y-%m-%d %H:%M:%S")
+        formatter = BotFormatter()
+        color_formatter = BotFormatter(color=True)
 
         logfile = lambda f: path.join(log_dir, f)
 
+        if not path.isdir(log_dir):
+            if not path.exists(log_dir):
+                mkdir(log_dir, 0700)
+            else:
+                msg = "log_dir ({0}) exists but is not a directory!"
+                print msg.format(log_dir)
+                exit(1)
+
         main_handler = hand(logfile("bot.log"), "midnight", 1, 7)
         error_handler = hand(logfile("error.log"), "W6", 1, 4)
         debug_handler = hand(logfile("debug.log"), "H", 1, 6)
-        stream_handler = logging.StreamHandler()
 
         main_handler.setLevel(logging.INFO)
         error_handler.setLevel(logging.WARNING)
         debug_handler.setLevel(logging.DEBUG)
-        stream_handler.setLevel(logging.DEBUG)
 
-        handlers = (main_handler, error_handler, debug_handler, stream_handler)
-        for h in handlers:
+        for h in (main_handler, error_handler, debug_handler):
             h.setFormatter(formatter)
             logger.addHandler(h)
 
+        stream_handler = logging.StreamHandler()
+        stream_handler.setLevel(logging.DEBUG)
+        stream_handler.setFormatter(color_formatter)
+        logger.addHandler(stream_handler)
+
     else:
         logger.addHandler(logging.NullHandler())
 
@@ -134,7 +144,7 @@ def load():
 
 def decrypt(key):
     """Use the key to decrypt passwords in our config file.
-    
+
     Call this if load() returns True. Catch password decryption errors and
     report them to the user.
     """
@@ -189,3 +199,30 @@ def schedule(minute, hour, month_day, month, week_day):
                 pass
 
     return tasks
+
+
+class BotFormatter(logging.Formatter):
+    def __init__(self, color=False):
+        self._format = super(BotFormatter, self).format
+        if color:
+            fmt = "[%(asctime)s %(lvl)s] %(name)s: %(message)s"
+            self.format = lambda record: self._format(self.format_color(record))
+        else:
+            fmt = "[%(asctime)s %(levelname)-8s] %(name)s: %(message)s"
+            self.format = self._format
+        datefmt = "%Y-%m-%d %H:%M:%S"
+        super(BotFormatter, self).__init__(fmt=fmt, datefmt=datefmt)
+
+    def format_color(self, record):
+        l = record.levelname.ljust(8)
+        if record.levelno == logging.DEBUG:
+            record.lvl = l.join(("\x1b[37m", "\x1b[0m"))
+        if record.levelno == logging.INFO:
+            record.lvl = l.join(("\x1b[32m", "\x1b[0m"))
+        if record.levelno == logging.WARNING:
+            record.lvl = l.join(("\x1b[36m", "\x1b[0m"))
+        if record.levelno == logging.ERROR:
+            record.lvl = l.join(("\x1b[33m", "\x1b[0m"))
+        if record.levelno == logging.CRITICAL:
+            record.lvl = l.join(("\x1b[31m", "\x1b[0m"))
+        return record