#!/usr/bin/env python

import sys

from subprocess import call, check_call, CalledProcessError
from os import listdir, unlink
from os.path import join, isdir, islink, exists
from optparse import OptionParser
from glob import glob

from rainbow.util import trace, make_reporter

def active_uid(uid):
    cmd = ['/usr/bin/pgrep', '-U', uid]
    ret = call(cmd, stdout=open('/dev/null'))
    if ret in (0, 1):
        return ret == 0
    raise CalledProcessError(ret, cmd)

def sticky_uid(spool, uid):
    return exists(join(spool, "sticky_uids", uid))

def gc_uid(log, spool, uid):
    """This function conservatively attempts to garbage-collect stale uid
    reservations.
    """

    # XXX: D-Bus caches passwd-db data!
    # XXX: D-Bus uses a fixed 1k passwd buffer - be careful with long paths &
    #      comments

    reservation = join(spool, 'uid_pool', uid)
    assert not isdir(reservation) and not islink(reservation)

    # XXX: We perform several execv()'s as root based on strings derived from
    # this 'uid' parameter, which originates as a file-name in a user-writable
    # directory. Better ideas for input validation would be welcome.
    uid_num = int(uid)
    assert uid_num >= 1000 and uid_num <= 65534  # XXX: magic numbers from util/spool.py

    if active_uid(uid) or sticky_uid(spool, uid):
        log(1, "skipped uid %s", uid)
        return

    for table in ('uid_to_instance_dir', 'uid_to_home_dir', 'uid_to_gid', 'uid_to_xephyr_auth', 'uid_to_xephyr_cookie', 'uid_to_xephyr_display'):
        row = join(spool, table, uid)
        # NB: it is important that rm -rf doesn't follow links. <MS>
        cmd = ['/bin/rm', '-r', '-f', row]
        log(2, "%s", ' '.join(cmd))
        check_call(cmd)

    for row in glob(join(spool, 'gid_to_members', '*', uid)):
        # NB: it is important that rm -rf doesn't follow links. <MS>
        cmd = ['/bin/rm', '-r', '-f', row]
        log(2, "%s", ' '.join(cmd))
        check_call(cmd)

    # So long as we unlink the reservation last, we run no risk of seeing inconsistency
    unlink(reservation)
    log(1, "cleaned uid %s", uid)

def gc_spool(log, spool):
    ret = 0
    uspool = join(spool, 'uid_pool')
    if exists(uspool) and isdir(uspool):
        for maybe_uid in listdir(uspool):
            try: gc_uid(log, spool, maybe_uid)
            except KeyboardInterrupt:
                raise
            except:
                trace()
                ret = 1
    else:
        log(1, "Skipping spool %s", spool)
        ret = 1
    return ret

def main():
    sys.excepthook = trace

    parser = OptionParser(version='0.1')
    parser.add_option('-v', '--verbose', default=0, action='count',
                      help='Verbosity. Repeat for more verbose output.')
    parser.add_option('-q', '--quiet', default=False, action='store_true',
                      help='Quiet. Disable all output.')
    parser.add_option('-s', '--spool', default="/var/spool/rainbow/2",
                      help='Location of the rainbow spool.')

    opts, _ = parser.parse_args()

    report = make_reporter(opts.verbose, opts.quiet, sys.stdout)

    def check_spool(opts):
        assert exists(opts.spool) and isdir(opts.spool)
        return opts.spool

    spool = check_spool(opts)
    return gc_spool(report, spool)

if __name__ == "__main__":
    exit(main())