From 5bf6e9f5b6a247c906d51d86317755b99c3d731a Mon Sep 17 00:00:00 2001
Message-Id: <5bf6e9f5b6a247c906d51d86317755b99c3d731a.1716253983.git.mdw@distorted.org.uk>
From: Mark Wooding <mdw@chiark.greenend.org.uk>
Date: Sun, 24 May 2015 16:45:32 +0100
Subject: [PATCH] catacomb/pwsafe.py: Make `PW' be a context manager, and use
 it.
Organization: Straylight/Edgeware

From: Mark Wooding <mdw@distorted.org.uk>

---
 catacomb/pwsafe.py |  7 ++++
 pwsafe             | 80 ++++++++++++++++++++++------------------------
 2 files changed, 46 insertions(+), 41 deletions(-)

diff --git a/catacomb/pwsafe.py b/catacomb/pwsafe.py
index 1e8e4db..82f2a46 100644
--- a/catacomb/pwsafe.py
+++ b/catacomb/pwsafe.py
@@ -299,4 +299,11 @@ class PW (object):
       if k[0] == '$': yield me.unpack(me.db[k])[0]
       k = me.db.nextkey(k)
 
+  ## Context protocol.
+
+  def __enter__(me):
+    return me
+  def __exit__(me, excty, excval, exctb):
+    me.db.close()
+
 ###----- That's all, folks --------------------------------------------------
diff --git a/pwsafe b/pwsafe
index 8a51bd7..5f2e75d 100644
--- a/pwsafe
+++ b/pwsafe
@@ -89,65 +89,63 @@ def cmd_create(av):
 
 def cmd_changepp(av):
   if len(av) != 0: return 1
-  pw = PW(file, writep = True)
-  pw.changepp()
+  with PW(file, writep = True) as pw: pw.changepp()
 
 def cmd_find(av):
   if len(av) != 1: return 1
-  pw = PW(file)
-  try: print pw[av[0]]
-  except KeyError, exc: die("Password `%s' not found." % exc.args[0])
+  with PW(file) as pw:
+    try: print pw[av[0]]
+    except KeyError, exc: die("Password `%s' not found" % exc.args[0])
 
 def cmd_store(av):
-  if len(av) < 1 or len(av) > 2:
-    return 1
+  if len(av) < 1 or len(av) > 2: return 1
   tag = av[0]
-  pw = PW(file, writep = True)
-  if len(av) < 2:
-    pp = C.getpass("Enter passphrase `%s': " % tag)
-    vpp = C.getpass("Confirm passphrase `%s': " % tag)
-    if pp != vpp: die("passphrases don't match")
-  elif av[1] == '-':
-    pp = stdin.readline()
-  else:
-    pp = av[1]
-  pw[av[0]] = chomp(pp)
+  with PW(file, writep = True) as pw:
+    if len(av) < 2:
+      pp = C.getpass("Enter passphrase `%s': " % tag)
+      vpp = C.getpass("Confirm passphrase `%s': " % tag)
+      if pp != vpp: die("passphrases don't match")
+    elif av[1] == '-':
+      pp = stdin.readline()
+    else:
+      pp = av[1]
+    pw[av[0]] = chomp(pp)
 
 def cmd_copy(av):
   if len(av) < 1 or len(av) > 2: return 1
-  pw_in = PW(file)
-  pw_out = PW(av[0], 'w')
-  if len(av) >= 3: pat = av[1]
-  else: pat = None
-  for k in pw_in:
-    if pat is None or fnmatch(k, pat): pw_out[k] = pw_in[k]
+  with PW(file) as pw_in:
+    with PW(av[0], writep = True) as pw_out:
+      if len(av) >= 3: pat = av[1]
+      else: pat = None
+      for k in pw_in:
+        if pat is None or fnmatch(k, pat): pw_out[k] = pw_in[k]
 
 def cmd_list(av):
   if len(av) > 1: return 1
-  pw = PW(file)
-  if len(av) >= 1: pat = av[0]
-  else: pat = None
-  for k in pw:
-    if pat is None or fnmatch(k, pat): print k
+  with PW(file) as pw:
+    if len(av) >= 1: pat = av[0]
+    else: pat = None
+    for k in pw:
+      if pat is None or fnmatch(k, pat): print k
 
 def cmd_topixie(av):
   if len(av) > 2: return 1
-  pw = PW(file)
-  pix = C.Pixie()
-  if len(av) == 0:
-    for tag in pw: pix.set(tag, pw[tag])
-  else:
-    tag = av[0]
-    if len(av) >= 2: pptag = av[1]
-    else: pptag = av[0]
-    pix.set(pptag, pw[tag])
+  with PW(file) as pw:
+    pix = C.Pixie()
+    if len(av) == 0:
+      for tag in pw: pix.set(tag, pw[tag])
+    else:
+      tag = av[0]
+      if len(av) >= 2: pptag = av[1]
+      else: pptag = av[0]
+      pix.set(pptag, pw[tag])
 
 def cmd_del(av):
   if len(av) != 1: return 1
-  pw = PW(file, writep = True)
-  tag = av[0]
-  try: del pw[tag]
-  except KeyError, exc: die("Password `%s' not found." % exc.args[0])
+  with PW(file, writep = True) as pw:
+    tag = av[0]
+    try: del pw[tag]
+    except KeyError, exc: die("Password `%s' not found" % exc.args[0])
 
 commands = { 'create': [cmd_create,
                         '[-c CIPHER] [-h HASH] [-m MAC] [PP-TAG]'],
-- 
[mdw]