Minor re-factoring of lock script

1 files changed, 19 insertions, 11 deletions

diff --git a/lock-repo.py b/lock-repo.py
index f651d3a..e3de55a 100755
--- a/lock-repo.py
+++ b/lock-repo.py
@@ -11,26 +11,34 @@ Mercurial Shared SSH Repo Lock Script
 import os
 import repolib
 
-def main(argv):
-    log = repolib.get_logger('validate-login')
 
-    if ('SSH_HG_REPO' not in os.environ or
-            'SSH_HG_USER' not in os.environ):
-        log.error("Failed to execute pre-lock checks")
+def check_repo(env, log):
+    repo = repolib.Repository(env['SSH_HG_REPO'])
+    repo.load_from_hgrc()
+
+    if not repo.can_be_written_by(env['SSH_HG_USER']):
+        log.error("You can not write to this repository.")
         return 1
 
+    return 0
+
+
+def main(argv):
+    log = repolib.get_logger('validate-login')
+
     try:
-        repo = repolib.Repository(os.environ['SSH_HG_REPO'])
-        repo.load_from_hgrc()
+        return check_repo(os.environ, log)
+
     except IOError:
-        log.error("Could not load repository config")
+        log.error("Could not load repository config.")
+        log.error("Failed to execute pre-lock checks.")
         return 1
 
-    if not repo.can_be_written_by(os.environ['SSH_HG_USER']):
-        log.error("You can not write to this repository")
+    except KeyError:
+        log.error("Incomplete or corrupted environment.")
+        log.error("Failed to execute pre-lock checks.")
         return 1
 
-    return 0
 
 if __name__ == "__main__":
     import sys