# -----------------------------------------------------------------------
# This is a script to edit/add/delete environment variables in a process
# (c) Hex-Rays
#
import idaapi
import idc
from idaapi import Appcall, COLSTR

# ----------------------------------------------------------------------
class EnvUtil(object):
    """Environment variable management utility class"""

    def __init__(self):
        """Get an Appcall instances to the needed APIs"""
        self.GetEnvironmentStrings  = Appcall.proto("kernel32_GetEnvironmentStringsA", "unsigned int __stdcall GetEnvironmentStrings();")
        self.SetEnvironmentVariable = Appcall.proto("kernel32_SetEnvironmentVariableA", "int __stdcall SetEnvironmentVariable(const char *lpName, const char *lpValue);")
        self.FreeEnvironmentStrings = Appcall.proto("kernel32_FreeEnvironmentStringsA", "int __stdcall FreeEnvironmentStrings(unsigned int block);")

    def ParseVar(self, x):
        p = x.find('=')
        if p == 0:
            return (None, x)
        return (x[:p], x[p+1:])

    def Variables(self):
        """Returns all environment blocks"""

        # Get all environment strings
        env_ptr = self.GetEnvironmentStrings()
        if env_ptr == 0:
            return None

        # Always refresh the memory after a call that could change the memory configuration
        idc.RefreshDebuggerMemory()

        # Open process memory
        f = idaapi.loader_input_t()
        f.open_memory(env_ptr)

        # Parse all environment variables into a list of variables
        vars = []
        var = []
        while True:
            ch = f.get_char()
            if not ch:
                break
            if ch == '\x00':
                # double-null -> end of list
                if not var:
                    break
                vars.append(''.join(var))
                var = []
            else:
                var.append(ch)
        f.close()

        # Free the block
        self.FreeEnvironmentStrings(env_ptr)

        return vars

    def DelVar(self, var):
        """Deletes an environment variable by setting its value to NULL"""
        return self.SetEnvironmentVariable(var, 0)

    def EditVar(self, k, v):
        """Edits an environment variable"""
        return self.SetEnvironmentVariable(k, v)

# -----------------------------------------------------------------------
class envviewer_t(idaapi.simplecustviewer_t):
    def Create(self):
        title = "Environment variable editor"
        win = idaapi.find_tform(title)
        if win:
            idaapi.open_tform(win, 0)
            return True

        self.env = EnvUtil()

        # Create the customviewer
        if not idaapi.simplecustviewer_t.Create(self, title):
            return False

        self.id_refresh = self.AddPopupMenu("Refresh")
        self.AddPopupMenu("-")
        self.id_edit    = self.AddPopupMenu("Edit")
        self.id_delete  = self.AddPopupMenu("Delete")
        self.AddPopupMenu("-")
        self.id_insert  = self.AddPopupMenu("Insert")
        self.Populate()

        return True

    def GetSelVar(self):
        line = self.GetCurrentLine(notags = True)
        if not line:
            return (None, None)
        return self.env.ParseVar(line)

    def DelSelVar(self):
        k, v = self.GetSelVar()
        if not k:
            if not v:
                print "No selection"
            else:
                print "Read-only variable"
        if self.env.DelVar(k):
            print "Deleted -> (%s=%s)" % (k, v)
            self.Populate()

    def EditSelVar(self):
        k, v = self.GetSelVar()
        if not k:
            if not v:
                print "No selection"
            else:
                print "Read-only variable"
        v = idaapi.askstr(0, v, "Enter new value for %s variable:" % k)
        if v and self.env.EditVar(k, v):
            print "Edited -> (%s=%s)" % (k, v)
            self.Populate()

    def InsertVar(self):
        k = idaapi.askstr(0, "key", "Enter key:")
        if not k:
            return
        v = idaapi.askstr(idaapi.HIST_IDENT, "value", "Enter value:")
        if not v:
            return
        if self.env.EditVar(k, v):
            print "New variable created -> (%s=%s)" % (k, v)
            self.Populate()

    def Populate(self):
        """Populates the customviewer with the environment variables"""
        self.ClearLines()
        for var in self.env.Variables():
            k, v = self.env.ParseVar(var)
            if not k:
                self.AddLine(COLSTR(v, idaapi.SCOLOR_AUTOCMT))
            else:
                self.AddLine('%s=%s' % (COLSTR(k, idaapi.SCOLOR_IMPNAME), COLSTR(v, idaapi.SCOLOR_DSTR)))
        self.Refresh()

    def OnPopupMenu(self, menu_id):
        if menu_id == self.id_refresh:
            self.Populate()
        elif menu_id == self.id_edit:
            self.EditSelVar()
        elif menu_id == self.id_insert:
            self.InsertVar()
        elif menu_id == self.id_delete:
            self.DelSelVar()
        else:
            return False
        return True

    def OnKeydown(self, vkey, shift):
        # ESCAPE?
        if vkey == 27:
            self.Close()
        # VK_DELETE
        elif vkey == 46:
            self.DelSelVar()
        # VK_INSERT
        elif vkey == 45:
            self.InsertVar()
        elif vkey == ord('E'):
            self.EditSelVar()
        else:
            return False
        return True

try:
    v = envviewer_t()
    if not v.Create():
        print "Failed to create!"
    else:
        v.Show()
except Exception, e:
    print "Failed to initialize! Make sure the debugger is running and the application is suspended!\n", e