diff --git a/DontTouchMyFlash/DontTouchMyFlash.csproj b/DontTouchMyFlash/DontTouchMyFlash.csproj index 075aaf9..0110370 100644 --- a/DontTouchMyFlash/DontTouchMyFlash.csproj +++ b/DontTouchMyFlash/DontTouchMyFlash.csproj @@ -52,6 +52,7 @@ + Form diff --git a/DontTouchMyFlash/FileUtil.cs b/DontTouchMyFlash/FileUtil.cs new file mode 100644 index 0000000..9a1a113 --- /dev/null +++ b/DontTouchMyFlash/FileUtil.cs @@ -0,0 +1,142 @@ +using System.Runtime.InteropServices; +using System.Diagnostics; +using System; +using System.Collections.Generic; + +static public class FileUtil +{ + [StructLayout(LayoutKind.Sequential)] + struct RM_UNIQUE_PROCESS + { + public int dwProcessId; + public System.Runtime.InteropServices.ComTypes.FILETIME ProcessStartTime; + } + + const int RmRebootReasonNone = 0; + const int CCH_RM_MAX_APP_NAME = 255; + const int CCH_RM_MAX_SVC_NAME = 63; + + enum RM_APP_TYPE + { + RmUnknownApp = 0, + RmMainWindow = 1, + RmOtherWindow = 2, + RmService = 3, + RmExplorer = 4, + RmConsole = 5, + RmCritical = 1000 + } + + [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)] + struct RM_PROCESS_INFO + { + public RM_UNIQUE_PROCESS Process; + + [MarshalAs(UnmanagedType.ByValTStr, SizeConst = CCH_RM_MAX_APP_NAME + 1)] + public string strAppName; + + [MarshalAs(UnmanagedType.ByValTStr, SizeConst = CCH_RM_MAX_SVC_NAME + 1)] + public string strServiceShortName; + + public RM_APP_TYPE ApplicationType; + public uint AppStatus; + public uint TSSessionId; + [MarshalAs(UnmanagedType.Bool)] + public bool bRestartable; + } + + [DllImport("rstrtmgr.dll", CharSet = CharSet.Unicode)] + static extern int RmRegisterResources(uint pSessionHandle, + UInt32 nFiles, + string[] rgsFilenames, + UInt32 nApplications, + [In] RM_UNIQUE_PROCESS[] rgApplications, + UInt32 nServices, + string[] rgsServiceNames); + + [DllImport("rstrtmgr.dll", CharSet = CharSet.Auto)] + static extern int RmStartSession(out uint pSessionHandle, int dwSessionFlags, string strSessionKey); + + [DllImport("rstrtmgr.dll")] + static extern int RmEndSession(uint pSessionHandle); + + [DllImport("rstrtmgr.dll")] + static extern int RmGetList(uint dwSessionHandle, + out uint pnProcInfoNeeded, + ref uint pnProcInfo, + [In, Out] RM_PROCESS_INFO[] rgAffectedApps, + ref uint lpdwRebootReasons); + + /// + /// Find out what process(es) have a lock on the specified file. + /// + /// Path of the file. + /// Processes locking the file + /// See also: + /// http://msdn.microsoft.com/en-us/library/windows/desktop/aa373661(v=vs.85).aspx + /// http://wyupdate.googlecode.com/svn-history/r401/trunk/frmFilesInUse.cs (no copyright in code at time of viewing) + /// + /// + static public List WhoIsLocking(string path) + { + uint handle; + string key = Guid.NewGuid().ToString(); + List processes = new List(); + + int res = RmStartSession(out handle, 0, key); + if (res != 0) throw new Exception("Could not begin restart session. Unable to determine file locker."); + + try + { + const int ERROR_MORE_DATA = 234; + uint pnProcInfoNeeded = 0, + pnProcInfo = 0, + lpdwRebootReasons = RmRebootReasonNone; + + string[] resources = new string[] { path }; // Just checking on one resource. + + res = RmRegisterResources(handle, (uint)resources.Length, resources, 0, null, 0, null); + + if (res != 0) throw new Exception("Could not register resource."); + + //Note: there's a race condition here -- the first call to RmGetList() returns + // the total number of process. However, when we call RmGetList() again to get + // the actual processes this number may have increased. + res = RmGetList(handle, out pnProcInfoNeeded, ref pnProcInfo, null, ref lpdwRebootReasons); + + if (res == ERROR_MORE_DATA) + { + // Create an array to store the process results + RM_PROCESS_INFO[] processInfo = new RM_PROCESS_INFO[pnProcInfoNeeded]; + pnProcInfo = pnProcInfoNeeded; + + // Get the list + res = RmGetList(handle, out pnProcInfoNeeded, ref pnProcInfo, processInfo, ref lpdwRebootReasons); + if (res == 0) + { + processes = new List((int)pnProcInfo); + + // Enumerate all of the results and add them to the + // list to be returned + for (int i = 0; i < pnProcInfo; i++) + { + try + { + processes.Add(Process.GetProcessById(processInfo[i].Process.dwProcessId)); + } + // catch the error -- in case the process is no longer running + catch (ArgumentException) { } + } + } + else throw new Exception("Could not list processes locking resource."); + } + else if (res != 0) throw new Exception("Could not list processes locking resource. Failed to get size of result."); + } + finally + { + RmEndSession(handle); + } + + return processes; + } +} \ No newline at end of file diff --git a/DontTouchMyFlash/FlashPwner.cs b/DontTouchMyFlash/FlashPwner.cs index 3494873..6f5a8b3 100644 --- a/DontTouchMyFlash/FlashPwner.cs +++ b/DontTouchMyFlash/FlashPwner.cs @@ -106,10 +106,19 @@ namespace DontTouchMyFlash } } - public void PatchExe(string filepath) + public bool PatchExe(string filepath) { try { + Process[] lockingProcesses = FileUtil.WhoIsLocking(filepath).ToArray(); + foreach(Process proc in lockingProcesses) + { + DialogResult res = MessageBox.Show("Flash is currently in use by (" + proc.Id.ToString() + ")" + proc.ProcessName + "\nEnd Process?", "File in use :/", MessageBoxButtons.YesNo, MessageBoxIcon.Information); + if (res == DialogResult.Yes) + proc.Kill(); // DIE HHAHA + else + return true; + } byte[] fileData = File.ReadAllBytes(filepath); Int64 timestampLocation = GetPositionAfterMatch(fileData, Timestamp); @@ -125,10 +134,12 @@ namespace DontTouchMyFlash flashExes.Items.Remove(filepath); Application.DoEvents(); progressBar.Increment(1); + return false; } catch(Exception e) { MessageBox.Show(e.Message, "ERROR", MessageBoxButtons.OK, MessageBoxIcon.Error); + return true; } } @@ -164,11 +175,13 @@ namespace DontTouchMyFlash { copyFlashExes.Add(flashExe); } + bool errored = false; foreach (string flashExe in copyFlashExes) { - PatchExe(flashExe); + errored = PatchExe(flashExe); } - MessageBox.Show("Patched! Your flash should work again!!!", "SUCCESS", MessageBoxButtons.OK, MessageBoxIcon.Information); + if(!errored) + MessageBox.Show("Patched! Your flash should work again!!!", "SUCCESS", MessageBoxButtons.OK, MessageBoxIcon.Information); } else { @@ -180,18 +193,23 @@ namespace DontTouchMyFlash private void addFile_Click(object sender, EventArgs e) { OpenFileDialog ofd = new OpenFileDialog(); - ofd.Title = "PE Executable"; - ofd.Filter = "PE Executables (*.exe, *.ocx, *.dll)|*.dll;*.exe;*.ocx"; - ofd.ShowDialog(); - if (!CheckFileAndAdd(ofd.FileName)) + ofd.Title = "Flash Executable"; + ofd.Filter = "PE Executables (*.exe, *.ocx, *.dll)|*.dll;*.exe;*.ocx|ELF Executables (*.so, *.elf, *.dylib)|*.so;*.elf;*.dylib"; + DialogResult res = ofd.ShowDialog(); + if(res == DialogResult.OK) { - MessageBox.Show("File selected does not contain the killswitch timestamp, cannot patch!", "Timestamp Error", MessageBoxButtons.OK, MessageBoxIcon.Error); + if (!CheckFileAndAdd(ofd.FileName)) + { + MessageBox.Show("File selected does not contain the killswitch timestamp, cannot patch!", "Timestamp Error", MessageBoxButtons.OK, MessageBoxIcon.Error); + } } + } private void deleteFile_Click(object sender, EventArgs e) { - flashExes.Items.RemoveAt(flashExes.SelectedIndex); + if(flashExes.SelectedIndex >= 0) + flashExes.Items.RemoveAt(flashExes.SelectedIndex); } } } diff --git a/DontTouchMyFlash/app.manifest b/DontTouchMyFlash/app.manifest index 7306c02..47e8f8e 100644 --- a/DontTouchMyFlash/app.manifest +++ b/DontTouchMyFlash/app.manifest @@ -16,7 +16,7 @@ Remove this element if your application requires this virtualization for backwards compatibility. --> - +