Added lock file, last page dump for post-mortem diagnostics and more error handling

fetch/banking-ib.py

@@ -188,6 +188,7 @@ class IBFetcher(object):
     BASE = "https://secure.ideabank.pl/"
     START_DATE = "01.11.2016"
     def __init__(self):
+        self._soup = None
         self.token = None
         self.s = requests.Session()
         self.s.headers.update(
@@ -201,6 +202,15 @@ class IBFetcher(object):
              "Cache-Control": "no-cache"
                 })
 
+    def _makesoup(data):
+        self._soup = bs4.BeautifulSoup(data)
+        return self._soup
+    
+    def _dump():
+        fn = config["DUMP_FILE"]
+        print "[w] Dumping the last page to {}".format(fn)
+        open(fn, 'w').write(unicode(self._soup).encode('utf-8'))
+
     def _getraw(self, page, params = {}):
         url = self.BASE + page
         r = self.s.get(url, params=params)
@@ -212,7 +222,7 @@ class IBFetcher(object):
     def _get(self, page):
         r = self._getraw(page)
         self.s.headers.update({"Referer": r.url})
-        soup = bs4.BeautifulSoup(r.text)
+        soup = _makesoup(r.text)
         self._gettoken(soup)
         self._hitjstoken(soup)
         return soup
@@ -228,6 +238,7 @@ class IBFetcher(object):
         r = self.s.post(url, data)
         print "[i] POST {} -> {}".format(page, r.status_code)
         if r.status_code != 200:
+            self._dump()
             raise Exception("return code %i" % r.status_code)
         return r
 
@@ -237,9 +248,10 @@ class IBFetcher(object):
         mdata.update(data)
         r = self._postraw(page, mdata)
         if re.search("forbidden",r.text) is not None:
+            self._dump()
             raise Exception("Received \"forbidden3\" response. Bad token?")
         self.s.headers.update({"Referer": r.url})
-        soup = bs4.BeautifulSoup(r.text)
+        soup = _makesoup(r.text)
         self._gettoken(soup)
         self._hitjstoken(soup)
         return soup
@@ -301,6 +313,11 @@ class IBFetcher(object):
             account["id"] = account_id
             wallet["accounts"][account["number"]] = account
         
+        if len(wallet["accounts"]) == 0:
+            print "[e] Empty accounts list. Undetected failed login? Aborting."
+            self._dump()
+            sys.exit(4)
+            
         return wallet
 
     def login(self, username, password):
@@ -313,12 +330,36 @@ class IBFetcher(object):
         login2_page = self._post("main/index", data)
         self._wait(3)
         
-        password2 = login2_page.find("input", attrs={"name": "password2"})["value"]
         data = {}
-        data["log2"] = username
-        data["password"] = password
-        data["password2"] = password2
+        
+        password2_input = login2_page.find("input", attrs={"name": "password2"})
+        if password2_input is None:
+            print "[e] Masked password screen encountered - aborting"
+            sys.exit(4)
+            #data["log"] = username
+            #data["log2"] = ""
+            #part_inputs = login2_page.find_all("input", class_="password_parts_inputs")
+            #print "[i] Filling out {} characters".format(len(part_inputs))
+            #for input in part_inputs:
+                #_,pos = input["name"].split("_")
+                #data["pass_"+str(pos)] = password[int(pos)]
+        else:
+            print "[i] Regular password screen encountered"
+            data["log2"] = username
+            data["password2"] = password2_input["value"]
+            data["password"] = password
+
         wallet_page = self._post("main/index", data)
+        
+        if wallet_page.find("div", class_="login_form"):
+            print "[e] Login failed, aborting"
+            self._dump()
+            try:
+                print "[e] Possible reason: {}".format(','.join(wallet_page.find("ul", class_="error_list").stripped_strings))
+            except:
+                pass # screw it, we're fucked anyway
+            sys.exit(4)
+        
         self._wait(2)
         return self.process_wallet_page(wallet_page)
 
@@ -354,6 +395,28 @@ class IBFetcher(object):
 def usage():
     pass
 
+def lock():
+    fn = config["LOCK_FILE"]
+    if os.path.isfile(fn):
+        print "[e] Lock file {} exists, aborting".format(fn)
+        sys.exit(3)
+    print "[i] Setting up lock file {}".format(fn)
+    open(fn,'w').close()
+    if not os.path.isfile(fn):
+        print "[e] Lock file {} somehow does not exist, aborting".format(fn)
+        sys.exit(3)
+        
+def release():
+    fn = config["LOCK_FILE"]
+    print "[i] Removing lock file {}".format(fn)
+    if not os.path.isfile(fn):
+        print "[e] Lock file {} somehow does not exist, WTF?".format(fn)
+        sys.exit(3)
+    os.remove(fn)
+    if os.path.isfile(fn):
+        print "[e] Lock file {} somehow still exists, WTF?".format(fn)
+        sys.exit(3)
+
 if __name__ == "__main__":
     try:
         opts, args = getopt(sys.argv[1:], "hcl:", ["help", "cached", "load=", "print-schema"])
@@ -400,6 +463,8 @@ if __name__ == "__main__":
         else:
             assert False, "unhandled option"
     
+    lock()
+    
     if cached:
         print "[i] Cached run - will not connect to the bank"
         if len(load_files) > 0:
@@ -455,4 +520,5 @@ if __name__ == "__main__":
     session.commit()
     
     print "[i] Done: ", stats
+    release()
     #print f.create_report().read()