From be69f65b9fdbb7944b3724cf0ba1e55b9ba66a3e Mon Sep 17 00:00:00 2001
From: Fredrik Tolf <fredrik@dolda2000.com>
Date: Sat, 5 Jun 2021 16:35:04 +0200
Subject: [PATCH] Make session saving and lookup more reusable.

---
 fulbank/data.py | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 fulbank/fsb.py  | 11 +---------
 netbank         | 48 ++++++++++++++++++--------------------------
 3 files changed, 81 insertions(+), 40 deletions(-)

diff --git a/fulbank/data.py b/fulbank/data.py
index 36efe20..ea73b8d 100644
--- a/fulbank/data.py
+++ b/fulbank/data.py
@@ -1,4 +1,4 @@
-import hashlib
+import os, pwd, hashlib, pickle
 
 def _localname(type):
     mod = type.__module__
@@ -46,3 +46,63 @@ class transaction(object):
 
     def __repr__(self):
         return "#<%s %s: %r>" % (_localname(type(self)), self.value, self.message)
+
+class session(object):
+    def save(self, filename):
+        with open(filename, "wb") as fp:
+            pickle.dump(self, fp)
+
+    @staticmethod
+    def load(filename):
+        with open(filename, "rb") as fp:
+            return pickle.load(fp)
+
+def getsessnam(name):
+    if name == "fsb":
+        from . import fsb
+        return fsb.session
+    raise ValueError("no such known session type: " + name)
+
+def _sesspath(name):
+    return os.path.join(pwd.getpwuid(os.getuid()).pw_dir, ".cache/fulbank", name)
+
+def defaultsess():
+    ret = os.getenv("NETBANKSESS")
+    if ret:
+        return ret
+    return "master"
+
+def loadsess(name=None, default=FileNotFoundError):
+    if name is None: name = defaultsess()
+    path = _sesspath(name)
+    if not os.path.exists(path):
+        if default is FileNotFoundError:
+            raise FileNotFoundError(name)
+        return default
+    return session.load(path)
+
+def savesess(sess, name=None):
+    if name is None: name = defaultsess()
+    path = _sesspath(name)
+    if sess is not None:
+        sessdir = os.path.dirname(path)
+        if not os.path.isdir(sessdir):
+            os.makedirs(sessdir)
+        return sess.save(_sesspath(name))
+    else:
+        if os.path.exists(path):
+            os.unlink(path)
+
+class savedsess(object):
+    def __init__(self, name=None):
+        if name is None: name = defaultsess()
+        self.name = name
+        self.sess = None
+
+    def __enter__(self):
+        self.sess = loadsess(self.name)
+        return self.sess
+
+    def __exit__(self):
+        savesess(self.sess, name)
+        self.sess = None
diff --git a/fulbank/fsb.py b/fulbank/fsb.py
index 166b5f5..7a4a4a8 100644
--- a/fulbank/fsb.py
+++ b/fulbank/fsb.py
@@ -171,7 +171,7 @@ class cardaccount(data.cardaccount):
                 yield cardtransaction(self, tx)
             page += 1
 
-class session(object):
+class session(data.session):
     def __init__(self, dsid):
         self.dsid = dsid
         self.auth = base64((serviceid + ":" + str(int(time.time() * 1000))).encode("ascii"))
@@ -358,12 +358,3 @@ class session(object):
     @classmethod
     def create(cls):
         return cls(getdsid())
-
-    def save(self, filename):
-        with open(filename, "wb") as fp:
-            pickle.dump(self, fp)
-
-    @classmethod
-    def load(cls, filename):
-        with open(filename, "rb") as fp:
-            return pickle.load(fp)
diff --git a/netbank b/netbank
index 5e7c474..a1520ca 100755
--- a/netbank
+++ b/netbank
@@ -1,9 +1,9 @@
 #!/usr/bin/python3
 
 import sys, os, getopt, pwd, operator
-from fulbank import auth
+from fulbank import auth, data
 
-sesstype = None
+sessname = data.defaultsess()
 sess = None
 
 def pfxmatch(pfx, item):
@@ -40,7 +40,7 @@ def find(seq, *, item=None, test=None, match=None, key=None, default=LookupError
         return default
 
 def usage(out):
-    out.write("usage: netbank [-h] BANK-ID COMMAND [ARGS...]\n")
+    out.write("usage: netbank [-h] [-s SESSION-ID] COMMAND [ARGS...]\n")
 
 def requiresess(fn):
     def wrap(cmd, args):
@@ -54,10 +54,15 @@ commands = {}
 
 def cmd_login(cmd, args):
     global sess
+    opts, args = getopt.getopt(args, "t:")
+    typename = sessname
+    for o, a in opts:
+        if o == "-t":
+            typename = a
     if len(args) < 1:
-        sys.stderr.write("usage: login TYPE\n")
+        sys.stderr.write("usage: login [-t BANK-ID] TYPE\n")
         sys.exit(1)
-    sess = sesstype.create()
+    sess = data.getsessnam(typename).create()
     if args[0] == "bankid":
         authfun = sess.auth_bankid
     elif args[0] == "token":
@@ -118,44 +123,29 @@ def cmd_lstxn(cmd, args):
 commands["lstxn"] = cmd_lstxn
 
 def main():
-    global sess, sesstype
+    global sess, sessname
 
-    opts, args = getopt.getopt(sys.argv[1:], "h")
+    opts, args = getopt.getopt(sys.argv[1:], "hs:")
     for o, a in opts:
         if o == "-h":
             usage(sys.stdout)
             sys.exit(0)
-    if len(args) < 2:
+        if o == "-s":
+            sessname = a
+    if len(args) < 1:
         usage(sys.stderr)
         sys.exit(1)
 
-    if args[0] == "fsb":
-        import fulbank.fsb
-        sesstype = fulbank.fsb.session
-    else:
-        sys.stderr.write("netbank: %s: unknown bank id\n" % (args[0]))
-        sys.exit(1)
-    sesspath = os.path.join(pwd.getpwuid(os.getuid()).pw_dir, ".cache/fulbank", args[0])
-    cmd = args[1]
-    args = args[2:]
+    cmd = args[0]
+    args = args[1:]
 
-    if os.path.exists(sesspath):
-        sess = sesstype.load(sesspath)
-    else:
-        sess = None
+    sess = data.loadsess(sessname, None)
     if cmd in commands:
         commands[cmd](cmd, args)
     else:
         sys.stderr.write("netbank: %s: unknown command\n" % (cmd))
         sys.exit(1)
-    if sess is not None:
-        sessdir = os.path.dirname(sesspath)
-        if not os.path.isdir(sessdir):
-            os.makedirs(sessdir)
-        sess.save(sesspath)
-    else:
-        if os.path.exists(sesspath):
-            os.unlink(sesspath)
+    data.savesess(sess, sessname)
 
 try:
     if __name__ == "__main__":
-- 
2.11.0