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        = [],

    50

    )

    51

    52

    return general

    53

    54

def options (**options) :

    55

    """

    56

        Synthensise options.

    57

    """

    58

    59

    return optparse.Values(options)

    60

    61

def apply_setid (options, rootok=None) :

    62

    """

    63

        Drop privileges if running as root.

    64

    65

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

    66

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

    67

    """

    68

    69

    # --uid -> pw

    70

    if not options.uid :

    71

        pw = None

    72

    elif options.uid.isdigit() :

    73

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

    74

    else :

    75

        pw = pwd.getpwnam(options.uid)

    76

    77

    # --gid -> gr

    78

    if not options.gid and not pw :

    79

        gr = None

    80

    elif not options.gid :

    81

        gr = grp.getgrgid(pw.pw_gid)

    82

    elif options.gid.isdigit() :

    83

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

    84

    else :

    85

        gr = grp.getgrnam(options.gid)

    86

    87

    if gr :

    88

        # XXX: secondary groups? seem to get cleared

    89

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

    90

        os.setgid(gr.gr_gid)

    91

    92

    if pw :

    93

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

    94

        os.setuid(pw.pw_uid)

    95

    96

    elif os.getuid() == 0 :

    97

        if rootok :

    98

            log.info("running as root")

    99

        else :

   100

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

   101

            sys.exit(2)

   102

   103

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

   104

    """

   105

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

   106

    """

   107

   108

    if not path or path == '-' :

   109

        # use stdin/out based on mode

   110

        stream, func = {

   111

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

   112

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

   113

        }[mode[0]]

   114

   115

        if charset :

   116

            return func(charset)(stream)

   117

        else :

   118

            return stream

   119

   120

    else :

   121

        if charset :

   122

            return codecs.open(path, mode, charset)

   123

        else :

   124

            return open(path, mode)

   125

   126

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

   127

    """

   128

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

   129

    """

   130

   131

    if paths :

   132

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

   133

    else:

   134

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

   135

   136

import configobj

   137

import copy

   138

import optparse

   139

   140

class Options (object) :

   141

    """

   142

        Custom optparse.Values implementation.

   143

   144

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

   145

            setattr(values, dest, ...)

   146

            values.ensure_value(dest, ...)

   147

    """

   148

   149

    def __init__ (self, defaults={ }) :

   150

        self._defaults  = defaults

   151

        self._options   = { }

   152

   153

    def __setattr__ (self, name, value) :

   154

        if name.startswith('_') :

   155

            self.__dict__[name] = value

   156

        else :

   157

            self._options[name] = value

   158

   159

    def ensure_value (self, name, default) :

   160

        if name in self._options :

   161

            pass

   162

        elif name in self._defaults :

   163

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

   164

        else :

   165

            self._options[name] = default

   166

   167

        return self._options[name]

   168

   169

    def _merge (self, options) :

   170

        """

   171

            Merge in options from given Options.

   172

        """

   173

   174

        # TODO: lists?

   175

        self._options.update(options._options)

   176

   177

    def __getattr__ (self, name) :

   178

        if name.startswith('_') :

   179

            raise AttributeError(name)

   180

   181

        if name in self._options :

   182

            return self._options[name]

   183

        elif name in self._defaults :

   184

            return self._defaults[name]

   185

        else :

   186

            raise AttributeError(name)

   187

   188

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

   189

    """

   190

        Load options from config.

   191

    """

   192

   193

    import configobj

   194

   195

    config = configobj.ConfigObj(config,

   196

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

   197

    )

   198

    config_options = Options(options._defaults)

   199

   200

    # load scalars

   201

    for scalar in config.scalars :

   202

        # option from config

   203

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

   204

   205

        if not option :

   206

            raise optparse.BadOptionError(scalar)

   207

   208

        # value from config

   209

        if option.takes_value() :

   210

            value = config[scalar]

   211

   212

        else :

   213

            # ignore

   214

            value = None

   215

   216

        # apply

   217

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

   218

   219

    # apply in actual options

   220

    config_options._merge(options)

   221

   222

    return config_options

   223

   224

def parse (parser, argv) :

   225

    """

   226

        Parse options, args from argv.

   227

    """

   228

   229

    options, args = parser.parse_args(argv[1:], values=Options(parser.defaults))

   230

   231

    if options.config :

   232

        options = apply_config(options, parser, options.config)

   233

   234

    return options, args

   235

   236

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

   237

    """

   238

        Apply the optparse options.

   239

    """

   240

   241

    if logname :

   242

        prefix = options.logname + ': '

   243

    else :

   244

        prefix = ''

   245

   246

    # configure

   247

    logging.basicConfig(

   248

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

   249

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

   250

        level       = options.loglevel,

   251

        filename    = options.log_file,

   252

    )

   253

   254

    # TODO: use --quiet for stdout output?

   255

    options.quiet = options.loglevel > logging.WARN

   256

   257

    if options.setuid :

   258

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

   259

            # set uid/gid

   260

            apply_setid(options, rootok=rootok)

   261

   262

    # enable debugging for specific targets

   263

    for logger in options.debug_module :

   264

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

   265

   266

def main (main) :

   267

    """

   268

        Run given main func.

   269

    """

   270

   271

    sys.exit(main(sys.argv))