Send user agent.

[plomlombot-irc.git] / plomlombot.py

diff --git a/plomlombot.py b/plomlombot.py
index b7829d2b9992ec230d5602772b1959641a52d0fd..28cc9ddbfc82aaf489fb7e76227ff7f91948ddc3 100755 (executable)
--- a/plomlombot.py
+++ b/plomlombot.py
@@ -11,6 +11,7 @@ import bs4
 import random
 import hashlib
 import os
+import signal
 import plomsearch
 import irclog
 
@@ -24,6 +25,12 @@ TWTFILE = ""
 DBDIR = os.path.expanduser("~/plomlombot_db")
 
 
+def write_to_file(path, mode, text):
+    f = open(path, mode)
+    f.write(text)
+    f.close()
+
+
 class ExceptionForRestart(Exception):
     pass
 
@@ -123,12 +130,9 @@ def handle_command(command, argument, notice, target, session):
 
     def addquote():
         if not os.access(session.quotesfile, os.F_OK):
-            quotesfile = open(session.quotesfile, "w")
-            quotesfile.write("QUOTES FOR " + target + ":\n")
-            quotesfile.close()
-        quotesfile = open(session.quotesfile, "a")
-        quotesfile.write(argument + "\n")
-        quotesfile.close()
+            write_to_file(session.quotesfile, "w",
+                          "QUOTES FOR " + target + ":\n")
+        write_to_file(session.quotesfile, "a", argument + "\n")
         quotesfile = open(session.quotesfile, "r")
         lines = quotesfile.readlines()
         quotesfile.close()
@@ -177,9 +181,11 @@ def handle_command(command, argument, notice, target, session):
             if len(results) == 0:
                 notice("NO QUOTES MATCHING QUERY")
             else:
-                for result in results:
-                    notice("QUOTE #" + str(result[0] + 1) + " : "
-                           + result[1][-1])
+                if len(results) > 3:
+                    notice("SHOWING 3 OF " + str(len(results)) + " QUOTES")
+                for result in results[:3]:
+                    notice("QUOTE #" + str(result[0] + 1) + ": "
+                           + result[1][:-1])
             return
         else:
             i = random.randrange(len(lines))
@@ -333,18 +339,34 @@ def handle_url(url, notice, show_url=False):
             handle_url(url, notice, True)
             return True
 
+    class TimeOut(Exception):
+        pass
+
+    def timeout_handler(ignore1, ignore2):
+        raise TimeOut("timeout")
+
+    signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(15)
     try:
-        r = requests.get(url, timeout=15)
+        r = requests.get(url, headers = {'User-Agent': 'plomlombot'}, stream=True)
+        r.raw.decode_content = True
+        text = r.raw.read(10000000+1)
+        if len(text) > 10000000:
+            raise ValueError('Too large a response')
     except (requests.exceptions.TooManyRedirects,
             requests.exceptions.ConnectionError,
             requests.exceptions.InvalidURL,
+            TimeOut,
             UnicodeError,
+            ValueError,
             requests.exceptions.InvalidSchema) as error:
+        signal.alarm(0)
         notice("TROUBLE FOLLOWING URL: " + str(error))
-        return
+        return False
+    signal.alarm(0)
     if mobile_twitter_hack(url):
-        return
-    title = bs4.BeautifulSoup(r.text, "html5lib").title
+        return True
+    title = bs4.BeautifulSoup(text, "html5lib").title
     if title and title.string:
         prefix = "PAGE TITLE: "
         if show_url:
@@ -352,6 +374,7 @@ def handle_url(url, notice, show_url=False):
         notice(prefix + title.string.strip())
     else:
         notice("PAGE HAS NO TITLE TAG")
+    return True
 
 
 class Session:
@@ -386,15 +409,13 @@ class Session:
                 line = Line(":" + self.nickname + "!~" + self.username +
                             "@localhost" + " " + line)
             now = datetime.datetime.utcnow()
-            logfile = open(self.rawlogdir + now.strftime("%Y-%m-%d") + ".txt", "a")
             form = "%Y-%m-%d %H:%M:%S UTC\t"
-            logfile.write(now.strftime(form) + " " + line.line + "\n")
-            logfile.close()
+            write_to_file(self.rawlogdir + now.strftime("%Y-%m-%d") + ".txt",
+                          "a", now.strftime(form) + " " + line.line + "\n")
             to_log = irclog.format_logline(line, self.channel)
             if to_log != None:
-                logfile = open(self.logdir + now.strftime("%Y-%m-%d") + ".txt", "a")
-                logfile.write(now.strftime(form) + " " + to_log + "\n")
-                logfile.close()
+                write_to_file(self.logdir + now.strftime("%Y-%m-%d") + ".txt",
+                              "a", now.strftime(form) + " " + to_log + "\n")
 
         def handle_privmsg(line):
 
@@ -408,17 +429,24 @@ class Session:
                 target = line.receiver
             msg = str.join(" ", line.tokens[3:])[1:]
             matches = re.findall("(https?://[^\s>]+)", msg)
+            url_count = 0
             for i in range(len(matches)):
-                handle_url(matches[i], notice)
+                if handle_url(matches[i], notice):
+                    url_count += 1
+                    if url_count == 3:
+                        notice("MAXIMUM NUMBER OF URLs TO PARSE PER MESSAGE "
+                               "REACHED")
+                        break
             if "!" == msg[0]:
                 tokens = msg[1:].split()
                 argument = str.join(" ", tokens[1:])
                 handle_command(tokens[0], argument, notice, target, self)
                 return
-            file = open(self.markovfile, "a")
-            file.write(msg + "\n")
-            file.close()
+            write_to_file(self.markovfile, "a", msg + "\n")
 
+        now = datetime.datetime.utcnow()
+        write_to_file(self.logdir + now.strftime("%Y-%m-%d") + ".txt", "a",
+                      "-----------------------\n")
         while True:
             if self.rmlogs > 0:
                 for f in os.listdir(self.logdir):