win_environment/path_append.py

from winreg import *
import argparse
import win32con
import win32gui
import re
import itertools

user = 'Environment'
system = r'SYSTEM\CurrentControlSet\Control\Session Manager\Environment'

class regkey(object):

    def __init__(self,root, path):
        self.path = path
        self.root = root
        reg = ConnectRegistry(None, root)
        self.key = OpenKey(reg, path, 0, KEY_ALL_ACCESS)

    def set_value(self, name, value,t=REG_EXPAND_SZ):
        SetValueEx(self.key, name, 0, t, value)
        win32gui.SendMessage(win32con.HWND_BROADCAST, win32con.WM_SETTINGCHANGE, 0, 'Environment')

    def get_value(self, name, ret_type = False):
        value, type_id = QueryValueEx(self.key, name)
        if ret_type:
            return value, type_id
        else:
            return value
            
class path_key(regkey):
    
    def __init__(self,root=HKEY_CURRENT_USER, path=user):
        super(path_key,self).__init__(root,path)
        self.path_set = self.get_path_setobj()
        
    def get_path_setobj(self):
        path = self.get_value('Path')
        path = set(path.split(';'))
        return path

    def format_path(self):
        path = self.path_set
        return ';'.join(sorted(path))

    def add(self, path):
        self.path_set.add(path)

    def remove(self, path):
        self.path_set.discard(path)

    def commit(self):
        path = self.format_path()
        self.set_value('Path', path)

    def check(self,path):
        return path in self.path_set

def search_remove(pathobj,pattern,action,regex=False,negate=False,dry_run=False):
    if negate:
        flt = itertools.filterfalse
    else:
        flt = filter
    if regex:
        pattern = re.compile(pattern,re.IGNORECASE)
        results = flt(lambda path: pattern.search(path),pathobj.path_set)
    else:
        pattern = pattern.lower()
        results = flt(lambda path: pattern in path.lower(),pathobj.path_set)
    for path in list(results):
        if dry_run:
            print('REM: {path}'.format(path=path))
        else:
            action(path)
            
if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog='path')
    
    parser.add_argument('-l','--list',action = 'store_true',help = 'list path entries')
    parser.add_argument('-s','--system',action = 'store_true',help = 'use system enviroment instead of current user')
    alts = parser.add_subparsers(dest='command')
    
    mod = alts.add_parser('mod')
    group = mod.add_mutually_exclusive_group(required=True)
    group.add_argument('-a', dest='add', action= 'store_true', help='add to PATH')
    group.add_argument('-r', dest='remove', action ='store_true', help='remove from PATH')
    mod.add_argument('-d',dest='dry_run',action='store_true',help='dry run')
    mod.add_argument('path', help = 'Path to be appended to PATH')
    
    rem = alts.add_parser('rem')
    rem.add_argument('--regex',action='store_true',help='use regex')
    rem.add_argument('-n',action = 'store_true',help = 'negate pattern')
    rem.add_argument('-d',dest='dry_run',action='store_true',help='dry run')
    rem.add_argument('pattern', help = 'pattern to search through path and remove matches')
    
    args = parser.parse_args()

    if args.system:
        try:
            path = path_key(HKEY_LOCAL_MACHINE,system)
        except PermissionError:
            print('Must be admin to view system keys')
            exit()
    else:
        path = path_key()
    
    if args.list:
        print('\n'.join(sorted(path.path_set)))
    elif args.command == 'mod':
        if args.dry_run:
            if args.add:
                print('ADD: {path}'.format(path=args.path))
            else:
                assert(path.check(args.path))
                print('REM: {path}'.format(path=args.path))
        else:
            if args.add:
                path.add(args.path)
            else:
                assert(path.check(args.path))
                path.remove(args.path)
            path.commit()
    elif args.command == 'map':
        search_remove(path,args.pattern,path.remove,args.regex,args.n,args.dry_run)
        if not args.dry_run:
            path.commit()
    else:
        parser.parse_args(['-h'])