#!/usr/bin/python

#
# Ceph - scalable distributed file system
#
# Copyright (C) 2011 New Dream Network
#
# This is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License version 2.1, as published by the Free Software
# Foundation.  See file COPYING.
#

"""
boto_tool.py: s3cmd-like tool for operating on s3

A lot of common s3 clients can't handle weird names.
But this little script can do it!
"""

from boto.s3.connection import OrdinaryCallingFormat
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from optparse import OptionParser
from sys import stderr
import ast
import ConfigParser
import boto
import os
import string
import sys

def list_all_buckets(conn, opts):
    blrs = conn.get_all_buckets()
    for b in blrs:
        if (opts.verbose):
            print b.__dict__
        else:
            print b

def list_objects(conn, opts):
    if opts.list_objects == True:
        prefix = None
    else:
        prefix = opts.list_objects
    bucket = conn.get_bucket(opts.bucket_name)
    for key in bucket.list(prefix = prefix):
        print key.name

def mkbucket(conn, opts):
    print "creating bucket '%s' ..." % opts.bucket_name
    bucket = conn.create_bucket(opts.bucket_name)
    print "done."
    return 0

def rmobjects(conn, opts):
    bucket = conn.get_bucket(opts.bucket_name)
    print "deleting all objects from bucket '%s' ..." % opts.bucket_name
    bucket = conn.get_bucket(opts.bucket_name)
    for key in bucket.list():
        print key.name
        bucket.delete_key(key)
    print "done."
    return 0

def rmbucket(conn, opts):
    bucket = conn.get_bucket(opts.bucket_name)
    print "deleting bucket '%s' ..." % opts.bucket_name
    bucket.delete()
    print "done."
    return 0

def get_bucket_acl(conn, opts):
    bucket = conn.get_bucket(opts.bucket_name)
    print "getting bucket acl..."
    xml = bucket.get_xml_acl()
    if (opts.getacl_file == "-"):
        print xml
    else:
        f = open(opts.getacl_file, "w")
        try:
            f.write(xml)
        finally:
            f.close()

def put_bucket_acl(conn, opts):
    bucket = conn.get_bucket(opts.bucket_name)
    print "setting bucket acl..."
    f = open(opts.putacl_file, "r")
    try:
        xml = f.read()
    finally:
        f.close()
    bucket.set_xml_acl(xml)
    return 0

def bucket_exists(conn, opts):
    bucket = conn.get_bucket(opts.bucket_name)
    if (bucket == None):
        print "bucket '%s' does not exist" % opts.bucket_name
        return 1
    else:
        print "found bucket '%s'." % opts.bucket_name
        if (opts.verbose):
            print bucket.__dict__
    return 0

def put_obj(conn, opts):
    if (opts.xattr == None):
        xattr = {}
    else:
        xattr = ast.literal_eval(opts.xattr)
        if not type(xattr) is dict:
          raise "You must supply a dictionary of xattr names to xattr " + \
            "values.  Example. \"{\"foo\" : \"bar\"}\""
    print "uploading to bucket: '%s', object name: '%s' \
xattr: '%s'" % (opts.bucket_name, opts.obj_name, str(xattr))
    bucket = conn.get_bucket(opts.bucket_name)
    k = Key(bucket)
    k.key = opts.obj_name
    k.metadata = xattr
    k.set_contents_from_filename(opts.put_file)
    return 0

def put_obj_acl(conn, opts):
    print "uploading object ACL to bucket: '%s', object name: '%s'" \
        % (opts.bucket_name, opts.obj_name)
    bucket = conn.get_bucket(opts.bucket_name)
    k = Key(bucket)
    k.key = opts.obj_name
    f = open(opts.putacl_file, "r")
    try:
        xml = f.read()
    finally:
        f.close()
    k.set_xml_acl(xml)
    return 0

def get_obj(conn, opts):
    print "downloading from bucket: '%s', object name: '%s'" % \
        (opts.bucket_name, opts.obj_name)
    bucket = conn.get_bucket(opts.bucket_name)
    k = Key(bucket)
    k.key = opts.obj_name
    if (opts.get_file == "-"):
        k.get_contents_to_file(sys.stdout)
    else:
        k.get_contents_to_filename(opts.get_file)
    return 0

def get_obj_acl(conn, opts):
    print "downloading object acl from bucket: '%s', object name: '%s'" % \
        (opts.bucket_name, opts.obj_name)
    bucket = conn.get_bucket(opts.bucket_name)
    k = Key(bucket)
    k.key = opts.obj_name
    xml = k.get_xml_acl()
    if (opts.getacl_file == "-"):
        print xml
    else:
        f = open(opts.getacl_file, "w")
        try:
            f.write(xml)
        finally:
            f.close()

def rm_obj(conn, opts):
    print "removing from bucket: '%s', object name: '%s'" \
        % (opts.bucket_name, opts.obj_name)
    bucket = conn.get_bucket(opts.bucket_name)
    bucket.delete_key(opts.obj_name)
    print "done."
    return 0

def head_obj(conn, opts):
    print "downloading from bucket: '%s', object name: '%s'" \
        % (opts.bucket_name, opts.obj_name)
    bucket = conn.get_bucket(opts.bucket_name)
    k = bucket.get_key(opts.obj_name)
    print k

def get_obj_xattr(conn, opts):
    print "downloading object xattr. bucket: '%s', object name: '%s'" % \
        (opts.bucket_name, opts.obj_name)
    bucket = conn.get_bucket(opts.bucket_name)
    k = bucket.get_key(opts.obj_name)
    print k.metadata

def get_obj_xattr(conn, opts):
    print "downloading object xattr. bucket: '%s', object name: '%s'" % \
        (opts.bucket_name, opts.obj_name)
    bucket = conn.get_bucket(opts.bucket_name)
    k = bucket.get_key(opts.obj_name)
    print k.metadata

########################## main ##########################
host = None
aws_access_key_id = None
secret_key = None

USAGE = """
boto_tool.py is a simple S3 client that can be used for testing.
It uses libboto.

ENVIRONMENT VARIABLES
S3TEST_CONF:      if this is set, we'll get the host, access key, and secret
key from this file.
AKEY              AWS access key
SKEY              AWS secret access key"""

parser = OptionParser(USAGE)
parser.add_option("-b", "--bucket-name",
    dest="bucket_name", help="specify bucket name")
parser.add_option("-l", "--list-objects", action="store_true", \
    dest="list_objects", help="list all objects")
parser.add_option("-L", "--list-objects-with-prefix",
    dest="list_objects", help="list all objects with the given prefix")
parser.add_option("--mkbucket", action="store_true",
    dest="mkbucket", help="create a bucket")
parser.add_option("--rmbucket", action="store_true",
    dest="rmbucket", help="remove a bucket")
parser.add_option("-o", "--object-name", dest="obj_name", help="object name")
parser.add_option("--put", dest="put_file", help="put FILE")
parser.add_option("--xattr", dest="xattr", help="xattrs to use")
parser.add_option("--get", dest="get_file", help="get to FILE")
parser.add_option("--rm", action="store_true", dest="rm", help="remove an object")
parser.add_option("--rmobjects", action="store_true", dest="rmobjects",
    help="remove all objects from a bucket")
parser.add_option("--rm_rf", action="store_true", dest="rm_rf",
    help="remove all objects from a bucket and remove the bucket")
parser.add_option("--head", action="store_true", dest="head",
    help="use the HEAD operation to find out about an object")
parser.add_option("--putacl", dest="putacl_file", help="set XML acl from FILE")
parser.add_option("--getacl", dest="getacl_file", help="dump XML acl into FILE")
parser.add_option("--get-obj-xattr", action="store_true",
    dest="get_obj_xattr", help="print the object's xattrs")
parser.add_option("-v", "--verbose", action="store_true", dest="verbose",
    help="be verbose")
(opts, args) = parser.parse_args()

if os.environ.has_key("S3TEST_CONF"):
    # use S3TEST_CONF
    path = os.environ["S3TEST_CONF"]
    cfg = ConfigParser.RawConfigParser()
    with file(path) as f:
        cfg.readfp(f)
        host = cfg.get("DEFAULT", "host")
        aws_access_key_id = cfg.get("s3 main", "access_key")
        secret_key = cfg.get("s3 main", "secret_key")
elif len(args) < 1:
    print >>stderr, "boto_tool: You must specify an action. Try --help."
    sys.exit(1)
else:
    # use environment variables and command line
    host = args[0]
    aws_access_key_id=os.environ["AKEY"]
    secret_key=os.environ["SKEY"]

conn = S3Connection(calling_format=OrdinaryCallingFormat(), is_secure=False,
                host = host,
                aws_access_key_id=aws_access_key_id,
                aws_secret_access_key=secret_key)

if not opts.bucket_name:
    sys.exit(list_all_buckets(conn, opts))
elif opts.list_objects:
    sys.exit(list_objects(conn, opts))
elif opts.mkbucket:
    sys.exit(mkbucket(conn, opts))
elif opts.rmobjects:
    sys.exit(rmobjects(conn, opts))
elif opts.rm_rf:
    ret =rmobjects(conn, opts)
    if (ret):
        sys.exit(ret)
    ret = rmbucket(conn, opts)
    sys.exit(ret)
elif opts.rmbucket:
    sys.exit(rmbucket(conn, opts))
elif not opts.obj_name:
    if opts.getacl_file:
        sys.exit(get_bucket_acl(conn, opts))
    elif opts.putacl_file:
        sys.exit(put_bucket_acl(conn, opts))
    else:
        sys.exit(bucket_exists(conn, opts))
elif opts.put_file:
    sys.exit(put_obj(conn, opts))
elif opts.get_file:
    sys.exit(get_obj(conn, opts))
elif opts.rm:
    sys.exit(rm_obj(conn, opts))
elif opts.head:
    sys.exit(head_obj(conn, opts))
elif opts.putacl_file:
    sys.exit(put_obj_acl(conn, opts))
elif opts.getacl_file:
    sys.exit(get_obj_acl(conn, opts))
elif opts.get_obj_xattr:
    sys.exit(get_obj_xattr(conn, opts))
elif opts.head:
    sys.exit(head_obj(conn, opts))
else:
    print "unknown arguments. Try --help"
    sys.exit(255)