Fix bug in import/export regarding some API params
[snf-image] / snf-image-host / pithcat
index b32e8ab..c0a0a60 100755 (executable)
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-# Copyright (C) 2011 GRNET S.A.
+# Copyright (C) 2011-2013 GRNET S.A.
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -25,44 +25,79 @@ Since the backend does not have a "root" account we use the account given in
 the URL as the user when connecting to the backend.
 """
 
-from optparse import OptionParser
-from sys import exit, stdout
+from optparse import OptionParser, OptionGroup
+from sys import exit, stdout, stderr
+from os import environ
+from binascii import hexlify, unhexlify
+from collections import namedtuple
 
-from pithos.backends.modular import ModularBackend
+try:
+    from pithos.backends.modular import ModularBackend
+except ImportError:
+    stderr.write("Pithos backend was not found.\n")
+    exit(2)
 
 
 parser = OptionParser(usage='%prog [options] <URL>')
-parser.add_option('--db', dest='db', metavar='URI',
-        help='SQLAlchemy URI of the database [REQUIRED]')
 parser.add_option('--data', dest='data', metavar='DIR',
-        help='path to the directory where data are stored [REQUIRED]')
+                  help='path to the directory where data are stored')
 parser.add_option('-s', action='store_true', dest='size', default=False,
-        help='print file size and exit')
-
-
-def urlsplit(url):
-    """Returns (accout, container, object) from a location string"""
-    
-    assert url.startswith('pithos://'), "Invalid URL"
-    t = url.split('/', 4)
-    assert len(t) == 5, "Invalid URL"
-    return t[2:5]
+                  help='print file size and exit')
+group = OptionGroup(
+    parser, "Dangerous Options",
+    "Caution: If the <URL> is a LocationURL (pithos://...), then you'll also "
+    "need to define a database URI. You can use the `--db' option to do so, "
+    "but this raises security concerns. For database URI's and pithos data "
+    "paths, the recommended way to define them is to use the PITHCAT_INPUT_DB "
+    "and PITHCAT_INPUT_DATA environmental variables respectfully.")
+group.add_option('--db', dest='db', metavar='URI',
+                 help='SQLAlchemy URI of the database', default=None)
+parser.add_option_group(group)
+
+LocationURL = namedtuple('LocationURL', ['account', 'container', 'object'])
+HashmapURL = namedtuple('HashmapURL', ['hash', 'size'])
+
+
+def parse_url(url):
+    if url.startswith('pithos://'):
+        t = url.split('/', 4)
+        assert len(t) == 5, "Invalid URL"
+        return LocationURL(*t[2:5])
+    elif url.startswith('pithosmap://'):
+        t = url.split('/', 3)
+        assert len(t) == 4, "Invalid URL"
+        return HashmapURL(*t[2:4])
+    else:
+        raise Exception("Invalid URL")
 
 
 def print_size(backend, url):
     """Writes object's size to stdout."""
-    
-    account, container, object = urlsplit(url)
-    meta = backend.get_object_meta(account, account, container, object, None)
-    print meta['bytes']
+    if type(url) is LocationURL:
+        account, container, object = url
+        meta = backend.get_object_meta(account, account, container, object,
+                                       None)
+        print meta['bytes']
+    elif type(url) is HashmapURL:
+        print url.size
+    else:
+        raise Exception("Invalid URL")
 
 
 def print_data(backend, url):
     """Writes object's size to stdout."""
-    
-    account, container, object = urlsplit(url)
-    size, hashmap = backend.get_object_hashmap(account, account, container,
-            object)
+
+    if type(url) is LocationURL:
+        account, container, object = url
+        size, hashmap = backend.get_object_hashmap(account, account, container,
+                                                   object)
+    elif type(url) is HashmapURL:
+        hashmap = [hexlify(x)
+                   for x in backend.store.map_get(unhexlify(url.hash))]
+        size = int(url.size)
+    else:
+        raise Exception("Invalid URL")
+
     for hash in hashmap:
         block = backend.get_block(hash)
         if len(block) > size:
@@ -73,13 +108,37 @@ def print_data(backend, url):
 
 def main():
     options, args = parser.parse_args()
-    if len(args) != 1 or not options.db or not options.data:
+    if len(args) != 1:
         parser.print_help()
         exit(1)
 
-    url = args[0]
-    backend = ModularBackend(None, options.db, None, options.data)
-    
+    url = parse_url(args[0])
+
+    if not options.data and 'PITHCAT_INPUT_DATA' not in environ:
+        stderr.write(
+            "Pithos data directory path is missing.\n"
+            "Use the PITHCAT_INPUT_DATA environmental variable (recommended) "
+            "or the --data command line option to define it.\n")
+        exit(1)
+
+    data_path = environ['PITHCAT_INPUT_DATA'] if not options.data else \
+        options.data
+
+    if options.db is None and 'PITHCAT_INPUT_DB' not in environ and \
+            type(url) is LocationURL:
+        stderr.write(
+            "Pithos database uri is missing.\n"
+            "Use the PITHCAT_INPUT_DB environmental variable (recommended) "
+            "or the --db command line option to define it.\n")
+        exit(1)
+
+    db_uri = environ['PITHCAT_INPUT_DB'] if not options.db else options.db
+
+    backend = ModularBackend(None,
+                             db_uri if type(url) is LocationURL else None,
+                             None,
+                             data_path)
+
     if options.size:
         print_size(backend, url)
     else:
@@ -87,3 +146,5 @@ def main():
 
 if __name__ == '__main__':
     main()
+
+# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :