Add download to browser command

browser.py

@@ -6,24 +6,20 @@ import prof
 import os
 import webbrowser
 import re
+import requests
+from os.path import expanduser
 
 _links = {}
 _lastlink = {}
 
 
-def _cmd_browser(url):
+def _get_url():
-    global _lastlink
     link = None
-
-    # use arg if supplied
-    if (url is not None):
-        link = url
-    else:
     jid = prof.get_current_recipient()
     room = prof.get_current_muc()
 
     # check if in chat window
-        if (jid is not None):
+    if jid is not None:
 
         # check for link from recipient
         if jid in _lastlink.keys():
@@ -32,18 +28,62 @@ def _cmd_browser(url):
             prof.cons_show("No links found from " + jid)
 
     # check if in muc window
-        elif (room is not None):
+    elif room is not None:
         if room in _lastlink.keys():
             link = _lastlink[room]
         else:
-                prof.cons_show("No links found from " + room)
+            prof.cons_show("No links found in " + room)
 
     # not in chat/muc window
     else:
         prof.cons_show("You must supply a URL to the /browser command")
 
-    # open the browser if link found
+    return link
-    if (link is not None):
+
+def _cmd_browser(arg1=None, arg2=None):
+    global _lastlink
+    link = None
+
+    if arg1 == "setdir":
+        if not arg2:
+            prof.cons_bad_cmd_usage("/browser")
+        else:
+            prof.settings_set_string("browser", "download.path", arg2)
+            prof.cons_show("browser.py set download path to" + arg2)
+
+        return
+
+    if arg1 == "get":
+        link = arg2 if arg2 is not None else _get_url()
+
+        if not link:
+            prof.log_debug("browser.py: no link found for get command")
+            return
+
+        folder = prof.settings_get_string("browser", "download.path", "")
+        if folder == "":
+            prof.cons_show("Download path not set, see /help browser")
+            return
+
+        if folder[0] == '~':
+            folder = expanduser("~") + folder[1:]
+
+        response = requests.get(link)
+        if response.status_code != 200:
+            prof.log_debug("browser.py: received non 200 response for get command")
+            return
+
+        filename = link.split("/")[-1]
+        full_path = folder + "/" + filename
+        prof.log_debug("browser.py: Saving link {link} to file {file}".format(link=link, file=full_path))
+        with open(full_path, 'wb') as f:
+            for chunk in response.iter_content(8):
+                f.write(chunk)
+
+        return
+
+    link = arg1 if arg1 is not None else _get_url();
+    if link is not None:
         prof.cons_show("Opening " + link + " in browser")
         _open_browser(link)
 
@@ -64,18 +104,28 @@ def _open_browser(url):
 def prof_init(version, status, account_name, fulljid):
     synopsis = [ 
         "/browser",
-        "/browser <url>"
+        "/browser <url>",
+        "/browser get <url>",
+        "/browser setdir <folder>"
     ]
-    description = "View a URL in the systems default browser. If no argument is supplied, the last URL in the current chat or room will be used. " \
+    description = (
+        "View a URL in the systems default browser, or download and view link contents " + 
+        "If no argument is supplied, the last URL in the current chat or room will be used. " + 
         "Tab autocompletion will go through all previous links in the current chat/room"
+    )
     args = [
-        [ "<url>", "URL to open in the browser" ]
+        [ "<url>",          "URL to open in the browser" ],
+        [ "download <url>", "Download contents of specified URL" ],
+        [ "setdir <dir>",   "Filesystem path to store downloaded content" ]
     ]
     examples = [
-        "/browser http://www.profanity.im"
+        "/browser http://www.profanity.im",
+        "/browser setdir ~/Downloads",
+        "/browser get http://www.movenoticias.com/wp-content/uploads/2015/05/Bruce-Iron-Maiden-850x583.jpg"
     ]
 
-    prof.register_command("/browser", 0, 1, synopsis, description, args, examples, _cmd_browser)
+    prof.register_command("/browser", 0, 2, synopsis, description, args, examples, _cmd_browser)
+    prof.completer_add("/browser", [ "get", "setdir" ])
 
 
 def prof_on_chat_win_focus(jid):
@@ -92,16 +142,20 @@ def prof_on_room_win_focus(room):
 
 def _process_message(jid, current_jid, message):
     links = re.findall(r'(https?://\S+)', message)
-    if (len(links) > 0):
+    if len(links) > 0:
         if jid not in _links:
             _links[jid] = []
+        
         # add to list of links for jid
         for link in links:
             if link not in _links[jid]:
+                prof.log_debug("browser.py: Saving {link} for {jid}".format(link=link, jid=jid))
                 _links[jid].append(link)
+        
         # add to autocompleter if message for current window
         if current_jid == jid:
             prof.completer_add("/browser", _links[jid])
+        
         # set last link for jid
         _lastlink[jid] = links[len(links)-1]