Update Wiretap Anonymous User & Group

update_wiretap_anon_usergroup.py

import os
import sys
import grp
import pwd
import re
import tempfile
import shutil
import getpass
import logging

WIRETAP_CFG = '/usr/discreet/wiretapgateway/cfg/wiretapgateway.cfg'
GROUP_KEY_NAME = "AnonymousGroup"
USER_KEY_NAME = "AnonymousUser"
DEFAULT_PATH = "/Volumes/Awesome-o"

# Setup a logger
log = logging.getLogger('my-logger')
log.setLevel(logging.WARNING)
hand = logging.StreamHandler()
form = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
hand.setFormatter(form)
log.addHandler(hand)

def get_group_name(path):
    """ Return the name of the group that owns a given path.

    Args:
        path (str): Valid file or directory path
    Returns:
        str. Name of group
    """
    gid = os.stat(path).st_gid
    group_name = grp.getgrgid(gid)[0]
    return group_name

def get_groups(user):
    """ Get a list of groups for a given user.
        http://stackoverflow.com/questions/9323834

    Args:
        user (str): Valid local user name.
    Returns:
        list. Contains all groups that the given user is part of.
    """
    gids = [g.gr_gid for g in grp.getgrall() if user in g.gr_mem]
    gid = pwd.getpwnam(user).pw_gid
    gids.append(grp.getgrgid(gid).gr_gid)
    return [grp.getgrgid(gid).gr_name for gid in gids]

def update_wiretap_cfg_key(key, value):
    """ Modifies the 'key' in WIRETAP_CFG to have the given 'value'.

    Args:
        key (str): Key name to find at the beginning of the cfg line.
        value (str): Value to place after the '='.
    """
    with open(WIRETAP_CFG, 'r') as fh:
        lines = fh.readlines()
    temp_file = tempfile.NamedTemporaryFile()
    for line in lines:
        regex = r"^\s*" + KEY_NAME + r"\s*=.*$"
        replacement = "%s=%s" % (KEY_NAME, value)
        temp_file.write(re.sub(regex, replacement, line))
    temp_file.flush()
    log.debug("Copying %s -> %s" % (temp_file.name, WIRETAP_CFG))
    shutil.copy(temp_file.name, WIRETAP_CFG)

def restart_wiretap():
    log.warning("Wiretap restart is not yet implemented")

def pick_directory(path):
    """ Provide a simple menu for selecting a subdirectory based on the input
        path.
        
    Args:
        path (str): Valid system path.
    Returns:
        str. new path based on user choice.
    """
    print "Currently selected directory: %s" % path
    dirs = [x for x in os.listdir(path) if os.path.isdir(os.path.join(path,x))]
    new_path = path
    if len(dirs)>0:
        def list_choices():
            print "Found %s subdirectories" % (len(dirs))
            print "  C. Use selected directory ('%s')" % path
            for i,d in enumerate(dirs):
                print "  %d. %s" % (i+1, d)
        list_choices()
        while True:
            value = raw_input("Select a path [C]: ").strip()
            if value.lower() in ['','c']:
                break
            try:
                choice = dirs[int(value)-1]
                new_path = os.path.join(path, choice)
                break
            except ValueError:
                pass
            except IndexError:
                log.error("Invalid choice.")
                list_choices()
                pass
    print "Using directory: %s" % new_path
    return new_path

def print_usage():
    print "Usage: %s <path>" % sys.argv[0]

def main():
    input_path = DEFAULT_PATH
    if len(sys.argv)==2:
        input_path = sys.argv[1]
    elif len(sys.argv)>2:
        print_usage()
        sys.exit(1)

    path = pick_directory(input_path)
    user = getpass.getuser()
    group_name = get_group_name(path)
    if group_name not in  get_groups(user):
        log.warning("'%s', is not part of group, '%s'.\n" % (user, group_name))

    update_wiretap_cfg_key(GROUP_KEY_NAME, group_name)
    update_wiretap_cfg_key(USER_KEY_NAME, user)

    restart_wiretap()

if __name__ == '__main__':
    main()