Generalise cli_auth_utils
[ep-engine.git] / management / cli_auth_utils.py
1 #!/usr/bin/env python
2
3 import clitool
4 import inspect
5 import mc_bin_client
6 import memcacheConstants
7 import sys
8 import os
9
10 def cmd_decorator(f):
11     """Decorate a function with code to authenticate based on 1-3
12     additional arguments."""
13
14     def g(*args, **kwargs):
15         mc = args[0]
16         spec = inspect.getargspec(f)
17         max = len(spec.args)
18         defaults = len(spec.defaults) if spec.defaults else 0
19         min = max - defaults
20
21         if len(args) < min:
22             print >> sys.stderr, ("Error: too few arguments - command "
23                                   "expected a minimum of %s but was passed "
24                                   "%s: %s"
25                                   % (min - 1, len(args) - 1, list(args[1:])))
26             sys.exit(2)
27
28         if spec.varargs is None:
29             if len(args) > max:
30                 print >> sys.stderr, ("Error: too many arguments - command "
31                                       "expected a maximum of %s but was passed "
32                                       "%s: %s"
33                                       % (max - 1, len(args) - 1, list(args[1:])))
34                 sys.exit(2)
35
36         bucket = kwargs.pop('bucketName', None) or 'default'
37         username = kwargs.pop('username', None) or bucket
38         password = kwargs.pop('password', None) or ''
39
40         if username:
41             try:
42                 mc.sasl_auth_plain(username, password)
43             except mc_bin_client.MemcachedError:
44                 print ("Authentication error for user:{0} bucket:{1}"
45                        .format(username, bucket))
46                 sys.exit(1)
47
48         mc.hello("{0} {1}".format(os.path.split(sys.argv[0])[1],
49                                 os.getenv("EP_ENGINE_VERSION",
50                                           "unknown version")))
51         try:
52             if kwargs.pop('allBuckets', None):
53                 buckets = mc.list_buckets()
54                 for bucket in buckets:
55                     print '*' * 78
56                     print bucket
57                     print
58                     mc.bucket_select(bucket)
59                     f(*args, **kwargs)
60             elif bucket is not None:
61                 mc.bucket_select(bucket)
62                 f(*args, **kwargs)
63             else:
64                 f(*args, **kwargs)
65         except mc_bin_client.ErrorEaccess:
66             print ("No access to bucket:{} - permission denied "
67                    "or bucket does not exist.".format(bucket))
68             sys.exit(1)
69
70     return g
71
72
73
74 def get_authed_clitool():
75     c = clitool.CliTool()
76
77     c.addFlag('-a', 'allBuckets', 'iterate over all buckets')
78     c.addOption('-b', 'bucketName', 'the bucket to get stats from (Default: default)')
79     c.addOption('-u', 'username', 'the user as which to authenticate (Default: bucketName)')
80     c.addOption('-p', 'password', 'the password for the bucket if one exists')
81
82     return c