PCM

pcm.py

#! /usr/bin/python

import argparse
import os
import pwd
import subprocess
import string
import random
import signal
import sys
import re

def fatrace_fork(login_name, pid_pcm_program, filename):
    pid = os.fork()
    if not pid:
        # Make this file non readable to other members
        so = se = open(filename, 'w')
        sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
        os.dup2(so.fileno(), sys.stdout.fileno())
        os.dup2(se.fileno(), sys.stderr.fileno())
        os.execvp("fatrace", ("fatrace",))
    else:
        print pid_pcm_program
        # We wait for the child process so 
        # has to get the correct status of the child 
        # program and not of su
        #print int(child_pid)
        os.waitpid(pid_pcm_program, 0)
        os.kill(pid, signal.SIGTERM)
        return
        # Clean up the temp file after extracing 
        # useful information from it 

def child(*args):
    # Get Pid first  and then use os.system
    # or use os.execve
    print args[0][0]
    arg_tuple = tuple(args[0])
    print args
    os.execvp(arg_tuple[0], arg_tuple)
    
def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for x in range(size))

# Maybe you can detect if the directory is a subvolume or not
def snapshot_file_path(path):
    path_split = path.split('/')
    if path_split[1] == 'home':
        path_split[1] = 'home/pcm'
    else:
        path_split[1] = 'pcm' + '/' + path_split[1]

    return '/'.join(path_split)



parser = argparse.ArgumentParser(description="Private Computing Mode Wrapper")
parser.add_argument("-p", "--policy", dest="filename", type=str, help="policy file for pcm")
parser.add_argument('program_args', type=str, nargs='+', help='Name of the program to run in pcm mode and its arguments')
parser.add_argument('-l', "--login", dest="login_name", type=str, help="who are you ?")
args = parser.parse_args()

# Sets for getting the files changed by 
# the pcm process and the other process
# in the directories mentioned in the policy file
discard_set = set()
discard_set_other = set()
askuser_set = set()
askuser_set_other = set()


if args.filename == None:
    parser.print_help()

print args.filename
print " ".join(args.program_args)
policy_dict = {}
regex_list = []
policy_values = ('askuser', 'keep', 'discard')
try:
    policy_file = open(args.filename, 'r')
    for line in policy_file:
        (filename, value) = line.split()
        value = value.strip()
        if filename[0] != '/' and filename[0] != '~':
            print 'Error in parsing the file path in the policy file'
            exit()
        if value not in policy_values:
            print 'pcm uses only askuser, keep, discard as policy values'
            exit()
        if filename in policy_dict:
            print 'Duplicate policies for the same file'
            exit()
        else:
            # Primitive shell expansion and some 
            # stuff for making regular expressions simpler
            if filename[0] == '~':
                filename = '/home/' + args.login_name + filename[1:]
            # This will make regular expression simpler some time in the 
            # future
            filename = filename.replace('*', '.*')
            print filename
            policy_dict[filename] = value
    print policy_dict

except IOError:
    print "Error reading policy file"
    exit()

# Change this to take the real user id
real_uid = 1000 #os.getresuid()[0]
login_name = pwd.getpwuid(os.getresuid()[0] )[0]
print login_name

# Before forking make a snapshot of the root 
# subvolume and the home subvolume 
# It assumes that you have a subvolume for /home
# and that is how it is suppose to be if its not you 
# doing it wrong

btrfs_create_root_snaphot = 'btrfs subvolume snapshot / /pcm'
btrfs_create_home_snaphot = 'btrfs subvolume snapshot /home /home/pcm'

btrfs_delete_root_snaphot = 'btrfs subvolume delete / /pcm'
btrfs_delete_home_snaphot = 'btrfs subvolume delete / /home/pcm'

subprocess.call(btrfs_create_root_snaphot.split())
subprocess.call(btrfs_create_home_snaphot.split())

# Fork 2 threads, 1 for actual program 
# and the other for the fatrace

pid = os.fork()
if not pid:
    os.setuid(real_uid)
    child(args.program_args)
else:
    filename = id_generator()
    filename = '/tmp/' + filename
    print filename
    fatrace_fork(login_name, pid, filename)
    try :
        pid = str(pid)
        f = open(filename, 'r')
    except IOError:
        print "Error opening the audit file"
        exit()


    for line in f:
        for key in policy_dict.keys():
            # regex for create as well
            search_regex = '.*\((.*)\):\s(CW|W)\s('  + key + ')'
                #print search_regex
                #print line
            m = re.search(search_regex, line)
                #print m
            if m is not None:
                if m.group(1) == pid:
                    if policy_dict[key] == 'discard':
                        discard_set.add(m.group(3))
                    else:
                        askuser_set.add(m.group(3))
                else:
                    if policy_dict[key] == 'discard':
                        discard_set_other.add(m.group(3))
                    else:
                            askuser_set_other.add(m.group(3))
                        
                print m.group(2)
                print m.group(3)
    print askuser_set
    print discard_set
    print askuser_set_other
    print discard_set_other
        
    askuser_set = askuser_set - askuser_set_other
    discard_set = discard_set - discard_set_other
       
    for discard_file in discard_set:
        try:
            snapshot_file = snapshot_file_path(discard_file)
            os.remove(discard_file)
            shutil.copy2(snapshot_file, discard_file)

        except OSError, why:
            print why
        except IOError, why:
            # Don't do anything
            pass
                            
                
            
            # Snapshotting stuff over here
    for askuser_file in askuser_set:
        x = raw_input('Do you want to keep the file ' + askuser_file + '[y/n]')
        if x == 'n':
            try:
                snapshot_file = snapshot_file_path(discard_file)
                os.remove(askuser_file)
                shutil.copy2(snapshot_file, discard_file)
            except OSError, why:
                print why
            except IOError, why:
                # Don't do anything
                pass
        
    subprocess.call(btrfs_delete_root_snaphot.split())
    subprocess.call(btrfs_delete_home_snaphot.split())
    try:
        os.remove(filename)
    except OSError, why:
        print why
        exit()