1

"""

     2

    CLI argument handling; common stuff: logging

     3

"""

     4

     5

import codecs

     6

import grp

     7

import logging

     8

import optparse

     9

import os

    10

import pwd

    11

import sys

    12

    13

import logging; log = logging.getLogger('pvl.args')

    14

    15

def parser (parser, setuid=None) :

    16

    """

    17

        Return an optparse.OptionGroup.

    18

    """

    19

    20

    if setuid is None :

    21

        # autodetect: only if we will be capable of

    22

        # XXX: use linux capabilities?

    23

        setuid = (os.geteuid() == 0)

    24

    25

    general = optparse.OptionGroup(parser, "General options")

    26

    27

    general.add_option('-q', '--quiet',     dest='loglevel', action='store_const', const=logging.ERROR, help="Less output")

    28

    general.add_option('-v', '--verbose',   dest='loglevel', action='store_const', const=logging.INFO,  help="More output")

    29

    general.add_option('-D', '--debug',     dest='loglevel', action='store_const', const=logging.DEBUG, help="Even more output")

    30

    general.add_option('--log-file',                                                                    help="Log to file")

    31

    general.add_option('--debug-module',    action='append', metavar='MODULE', 

    32

            help="Enable logging for the given logger/module name")

    33

    35

            help="Read option defaults from config")

    36

    parser.add_option('--config-encoding',  metavar='CHARSET',  default='utf-8',

    37

            help="Unicode decoding for config file")

    38

    39

    40

    if setuid :

    41

        general.add_option('--uid',             help="Change uid")

    42

        general.add_option('--gid',             help="Change gid")

    43

    44

    # defaults

    45

    parser.set_defaults(

    46

        setuid              = setuid,

    47

        logname             = parser.prog,

    48

        loglevel            = logging.WARN,

    49

        debug_module        = [],

    51

    )

    52

    53

    return general

    54

    55

def options (**options) :

    56

    """

    57

        Synthensise options.

    58

    """

    59

    60

    return optparse.Values(options)

    61

    62

def apply_setid (options, rootok=None) :

    63

    """

    64

        Drop privileges if running as root.

    65

    66

        XXX: this feature isn't very useful (import-time issues etc), but in certain cases (syslog-ng -> python),

    67

        it's difficult to avoid this without some extra wrapper tool..?

    68

    """

    69

    70

    # --uid -> pw

    71

    if not options.uid :

    72

        pw = None

    73

    elif options.uid.isdigit() :

    74

        pw = pwd.getpwuid(int(options.uid))

    75

    else :

    76

        pw = pwd.getpwnam(options.uid)

    77

    78

    # --gid -> gr

    79

    if not options.gid and not pw :

    80

        gr = None

    81

    elif not options.gid :

    82

        gr = grp.getgrgid(pw.pw_gid)

    83

    elif options.gid.isdigit() :

    84

        gr = grp.getgrgid(str(options.gid))

    85

    else :

    86

        gr = grp.getgrnam(options.gid)

    87

    88

    if gr :

    89

        # XXX: secondary groups? seem to get cleared

    90

        log.info("setgid: %s: %s", gr.gr_name, gr.gr_gid)

    91

        os.setgid(gr.gr_gid)

    92

    93

    if pw :

    94

        log.info("setuid: %s: %s", pw.pw_name, pw.pw_uid)

    95

        os.setuid(pw.pw_uid)

    96

    97

    elif os.getuid() == 0 :

    98

        if rootok :

    99

            log.info("running as root")

   100

        else :

   101

            log.error("refusing to run as root, use --uid 0 to override")

   102

            sys.exit(2)

   103

   104

def apply_file (path=None, mode='r', charset=None) :

   105

    """

   106

        Open (unicode-enabled) file from path, with - using stdio.

   107

    """

   108

   109

    if not path or path == '-' :

   110

        # use stdin/out based on mode

   111

        stream, func = {

   112

            'r':    (sys.stdin, codecs.getreader),

   113

            'w':    (sys.stdout, codecs.getwriter),

   114

        }[mode[0]]

   115

   116

        if charset :

   117

            return func(charset)(stream)

   118

        else :

   119

            return stream

   120

   121

    else :

   122

        if charset :

   123

            return codecs.open(path, mode, charset)

   124

        else :

   125

            return open(path, mode)

   126

   127

def apply_files (paths, *args, **opts) :

   128

    """

   129

        Open one or more files from given paths, defaulting to stdio.

   130

    """

   131

   132

    if paths :

   133

        return [apply_file(path, *args, **opts) for path in paths]

   134

    else:

   135

        return [apply_file(None, *args, **opts)]

   136

   137

import configobj

   138

import copy

   139

import optparse

   140

   141

class Options (object) :

   142

    """

   143

        Custom optparse.Values implementation.

   144

   145

        Passed to OptionParser.parse_args(), called by Option.take_action():

   146

            setattr(values, dest, ...)

   147

            values.ensure_value(dest, ...)

   148

    """

   149

   152

        self._defaults  = defaults

   153

        self._options   = { }

   154

   155

    def __setattr__ (self, name, value) :

   156

        if name.startswith('_') :

   157

            self.__dict__[name] = value

   158

        else :

   159

            self._options[name] = value

   160

   161

    def ensure_value (self, name, default) :

   162

        if name in self._options :

   163

            pass

   164

        elif name in self._defaults :

   165

            self._options[name] = copy.copy(self._defaults[name])

   166

        else :

   167

            self._options[name] = default

   168

   169

        return self._options[name]

   170

   171

    def _merge (self, options) :

   172

        """

   173

            Merge in options from given Options.

   174

        """

   175

   176

        # TODO: lists?

   177

        self._options.update(options._options)

   178

   179

    def __getattr__ (self, name) :

   180

        if name.startswith('_') :

   181

            raise AttributeError(name)

   182

   183

        if name in self._options :

   184

            return self._options[name]

   185

        elif name in self._defaults :

   186

            return self._defaults[name]

   187

        else :

   188

            raise AttributeError(name)

   189

   197

def apply_config (options, parser, config, encoding=None) :

   198

    """

   199

        Load options from config.

   200

    """

   201

   202

    import configobj

   203

   204

    config = configobj.ConfigObj(config,

   205

            encoding        = options.config_encoding if encoding is None else encoding,

   206

    )

   208

   209

    # load scalars

   210

    for scalar in config.scalars :

   211

        # option from config

   212

        option = parser._long_opt.get('--' + scalar)

   213

   214

        if not option :

   215

            raise optparse.BadOptionError(scalar)

   216

   217

        # value from config

   218

        if option.takes_value() :

   219

            value = config[scalar]

   220

   221

        else :

   222

            # ignore

   223

            value = None

   224

   225

        # apply

   226

        option.process(scalar, value, config_options, parser)

   227

   228

    # apply in actual options

   229

    config_options._merge(options)

   230

   231

    return config_options

   232

   233

def parse (parser, argv) :

   234

    """

   235

        Parse options, args from argv.

   236

    """

   237

   239

   242

   243

    return options, args

   244

   245

def apply (options, logname=None, rootok=True) :

   246

    """

   247

        Apply the optparse options.

   248

    """

   249

   250

    if logname :

   251

        prefix = options.logname + ': '

   252

    else :

   253

        prefix = ''

   254

   255

    # configure

   256

    logging.basicConfig(

   257

        # XXX: log Class.__init__ as Class, not __init__?

   258

        format      = prefix + '%(levelname)8s %(name)20s.%(funcName)s: %(message)s',

   259

        level       = options.loglevel,

   260

        filename    = options.log_file,

   261

    )

   262

   263

    # TODO: use --quiet for stdout output?

   264

    options.quiet = options.loglevel > logging.WARN

   265

   266

    if options.setuid :

   267

        if options.uid or options.gid or not rootok :

   268

            # set uid/gid

   269

            apply_setid(options, rootok=rootok)

   270

   271

    # enable debugging for specific targets

   272

    for logger in options.debug_module :

   273

        logging.getLogger(logger).setLevel(logging.DEBUG)

   274

   275

def main (main) :

   276

    """

   277

        Run given main func.

   278

    """

   279

   280

    sys.exit(main(sys.argv))