Kill process and delete associated files

kill_clean

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

"""Kill process and delete associated files

This script deletes all files in the system's tmp folder that have an
extension and PID in their name. By default it prompts the user for
confirmation before deleting the files.

Usage
-----

kill_clean [-h] [-u USER] [-f] [--interactive] [--quiet] PID
           [PID_ARGS ...]

positional arguments:
  PID                   Process to kill
  PID_ARGS              Arguments to pass to PID

optional arguments:
  -h, --help            show this help message and exit
  -u USER, --user USER  User who must own file to delete (regex)
  -f, --force           Do not prompt user whether for confirmation
  --interactive         Choose which files to delete
  --quiet               Do not print anything, just kill

Examples
--------

kill_clean 9999
kill_clean 9999 -f -9
"""

from __future__ import print_function

from os import path, linesep, unlink, system, stat
from tempfile import gettempdir
from getpass import getuser
from pwd import getpwuid
from glob import glob
import argparse
import re

# ---------------------------------------------------------------------
# CLI args

parser = argparse.ArgumentParser()
parser.add_argument('PID',
                    nargs   = 1,
                    type    = str,
                    metavar = 'PID',
                    help    = "Process to kill")
parser.add_argument('--user',
                    default  = getuser(),
                    dest     = 'USER',
                    type     = str,
                    help     = "User who must own file to delete (regex)",
                    required = False)
parser.add_argument('-f', '--force',
                    dest     = 'FORCE',
                    action   = 'store_true',
                    help     = "Do not prompt user whether for confirmation",
                    required = False)
parser.add_argument('--interactive',
                    dest     = 'INTERACTIVE',
                    action   = 'store_true',
                    help     = "Choose which files to delete",
                    required = False)
parser.add_argument('--quiet',
                    dest     = 'QUIET',
                    action   = 'store_true',
                    help     = "Do not print anything, just kill",
                    required = False)
cli_args, cli_kwargs = parser.parse_known_args()
cli_args = vars(cli_args)


if cli_args["QUIET"]:
    def print_verbose(*args, **kwargs):
        pass
else:
    def print_verbose(*args, **kwargs):
        print(*args, **kwargs)


# ---------------------------------------------------------------------
# Parse PID

pid        = cli_args['PID'][0]
pid_files_ = glob(path.join(gettempdir(), "*{0}*.*".format(pid)))
pid_msg    = "in '{0}' for PID {1}".format(gettempdir(), pid)

# If interactive, ask about individual files
# ------------------------------------------

sep = linesep + "\t"
if cli_args['INTERACTIVE']:
    print_verbose('Found ' + pid_msg)
    pid_files = []
    for pid_file in pid_files_:
        raw_msg = "\t{0}. Delete? Yes/No (Default: No) "
        delete  = raw_input(raw_msg.format(pid_file))
        if delete.lower() in ['y', 'ye', 'yes']:
            pid_files += [pid_file]
else:
    pid_files = pid_files_

# Files to delete
# ---------------

nfiles = len(pid_files)

if nfiles == 0:
    print_verbose('Nothing ' + pid_msg)
else:
    if cli_args['INTERACTIVE']:
        print_verbose()
        print_verbose(sep.join(['Files to delete ' + pid_msg] + pid_files))
    else:
        print_verbose(sep.join(['Found ' + pid_msg] + pid_files))

    # Prompt for confirmation w/o force; with interactive the default is yes
    # because we already asked the user.
    if cli_args['INTERACTIVE']:
        delete = raw_input("Yes/No? (Default: Yes) ")
        if delete.lower() in ['y', 'ye', 'yes', '']:
            rm = True
        else:
            rm = False
    elif cli_args['FORCE']:
        rm = True
    else:
        delete = raw_input("Delete? Yes/No (Default: No) ")
        if delete.lower() in ['y', 'ye', 'yes']:
            rm = True
        else:
            rm = False

    # If the user requested the files be deleted, do it. Else continue.
    if rm:
        pid_user    = cli_args['USER']
        pid_skipped = []
        for pid_file in pid_files:
            owner = getpwuid(stat(pid_file).st_uid).pw_name
            if re.match(pid_user, owner):
                unlink(pid_file)
            else:
                pid_skipped += [pid_file + ' (owned by {0})'.format(owner)]

        if len(pid_skipped) > 0:
            skip_msg = "Skipped files with an owner not matching '{0}':"
            print_verbose()
            print_verbose(sep.join([skip_msg.format(pid_user)] + pid_skipped))
    else:
        print_verbose("Nothing to delete.")

# Kill the process
# ----------------

pid_kill = 'kill {0} {1}'.format(' '.join(cli_kwargs), pid)
system(pid_kill)