You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

489 lines
19 KiB

4 years ago
4 years ago
  1. using System;
  2. using System.Collections.Generic;
  3. using System.Diagnostics;
  4. using System.Diagnostics.CodeAnalysis;
  5. using System.IO;
  6. using System.Linq;
  7. using System.Reflection;
  8. using System.Runtime.InteropServices;
  9. using System.Text.RegularExpressions;
  10. using System.Windows.Forms;
  11. using IPA.Patcher;
  12. namespace IPA
  13. {
  14. [SuppressMessage("ReSharper", "ClassNeverInstantiated.Global")]
  15. public class Program
  16. {
  17. [SuppressMessage("ReSharper", "InconsistentNaming")]
  18. public enum Architecture
  19. {
  20. x86,
  21. x64,
  22. Unknown
  23. }
  24. public const string FileVersion = "3.99.99.6";
  25. public static Version Version => Assembly.GetEntryAssembly().GetName().Version;
  26. public static readonly ArgumentFlag ArgHelp = new ArgumentFlag("--help", "-h") { DocString = "prints this message" };
  27. public static readonly ArgumentFlag ArgWaitFor = new ArgumentFlag("--waitfor", "-w") { DocString = "waits for the specified PID to exit", ValueString = "PID" };
  28. public static readonly ArgumentFlag ArgForce = new ArgumentFlag("--force", "-f") { DocString = "forces the operation to go through" };
  29. public static readonly ArgumentFlag ArgRevert = new ArgumentFlag("--revert", "-r") { DocString = "reverts the IPA installation" };
  30. public static readonly ArgumentFlag ArgNoRevert = new ArgumentFlag("--no-revert", "-R") { DocString = "prevents a normal installation from first reverting" };
  31. public static readonly ArgumentFlag ArgNoWait = new ArgumentFlag("--nowait", "-n") { DocString = "doesn't wait for user input after the operation" };
  32. public static readonly ArgumentFlag ArgStart = new ArgumentFlag("--start", "-s") { DocString = "uses the specified arguments to start the game after the patch/unpatch", ValueString = "ARGUMENTS" };
  33. public static readonly ArgumentFlag ArgLaunch = new ArgumentFlag("--launch", "-l") { DocString = "uses positional parameters as arguments to start the game after patch/unpatch" };
  34. [STAThread]
  35. public static void Main(string[] args)
  36. {
  37. Arguments.CmdLine.Flags(ArgHelp, ArgWaitFor, ArgForce, ArgRevert, ArgNoWait, ArgStart, ArgLaunch, ArgNoRevert).Process();
  38. if (ArgHelp)
  39. {
  40. Arguments.CmdLine.PrintHelp();
  41. return;
  42. }
  43. try
  44. {
  45. if (ArgWaitFor.HasValue)
  46. { // wait for process if necessary
  47. var pid = int.Parse(ArgWaitFor.Value);
  48. try
  49. { // wait for beat saber to exit (ensures we can modify the file)
  50. var parent = Process.GetProcessById(pid);
  51. Console.WriteLine($"Waiting for parent ({pid}) process to die...");
  52. parent.WaitForExit();
  53. }
  54. catch (Exception)
  55. {
  56. // ignored
  57. }
  58. }
  59. PatchContext context = null;
  60. Assembly AssemblyLibLoader(object source, ResolveEventArgs e)
  61. {
  62. // ReSharper disable AccessToModifiedClosure
  63. if (context == null) return null;
  64. var libsDir = context.LibsPathSrc;
  65. // ReSharper enable AccessToModifiedClosure
  66. var asmName = new AssemblyName(e.Name);
  67. var testFile = Path.Combine(libsDir, $"{asmName.Name}.{asmName.Version}.dll");
  68. if (File.Exists(testFile))
  69. return Assembly.LoadFile(testFile);
  70. Console.WriteLine($"Could not load library {asmName}");
  71. return null;
  72. }
  73. AppDomain.CurrentDomain.AssemblyResolve += AssemblyLibLoader;
  74. var argExeName = Arguments.CmdLine.PositionalArgs.FirstOrDefault(s => s.EndsWith(".exe"));
  75. if (argExeName == null)
  76. context = PatchContext.Create(new DirectoryInfo(Directory.GetCurrentDirectory()).GetFiles()
  77. .First(o => o.Extension == ".exe" && o.FullName != Assembly.GetEntryAssembly().Location)
  78. .FullName);
  79. else
  80. context = PatchContext.Create(argExeName);
  81. // Sanitizing
  82. Validate(context);
  83. if (ArgRevert || Keyboard.IsKeyDown(Keys.LMenu))
  84. Revert(context);
  85. else
  86. {
  87. Install(context);
  88. StartIfNeedBe(context);
  89. }
  90. }
  91. catch (Exception e)
  92. {
  93. Fail(e.Message);
  94. }
  95. WaitForEnd();
  96. }
  97. private static void WaitForEnd()
  98. {
  99. if (!ArgNoWait)
  100. {
  101. Console.ForegroundColor = ConsoleColor.DarkYellow;
  102. Console.WriteLine("[Press any key to continue]");
  103. Console.ResetColor();
  104. Console.ReadKey();
  105. }
  106. }
  107. private static void Validate(PatchContext c)
  108. {
  109. if (!Directory.Exists(c.DataPathDst) || !File.Exists(c.EngineFile))
  110. {
  111. Fail("Game does not seem to be a Unity project. Could not find the libraries to patch.");
  112. Console.WriteLine($"DataPath: {c.DataPathDst}");
  113. Console.WriteLine($"EngineFile: {c.EngineFile}");
  114. }
  115. }
  116. private static void Install(PatchContext context)
  117. {
  118. try
  119. {
  120. bool installFiles = true;
  121. // first, check currently installed version, if any
  122. if (File.Exists(Path.Combine(context.ProjectRoot, "winhttp.dll")))
  123. { // installed, so check version of installed assembly
  124. string injectorPath = Path.Combine(context.ManagedPath, "IPA.Injector.dll");
  125. if (File.Exists(injectorPath))
  126. {
  127. var verInfo = FileVersionInfo.GetVersionInfo(injectorPath);
  128. var fileVersion = new Version(verInfo.FileVersion);
  129. if (fileVersion > Version)
  130. installFiles = false;
  131. }
  132. }
  133. if (installFiles || ArgForce)
  134. {
  135. var backup = new BackupUnit(context);
  136. if (!ArgNoRevert)
  137. {
  138. Console.ForegroundColor = ConsoleColor.Cyan;
  139. Console.WriteLine("Restoring old version... ");
  140. if (BackupManager.HasBackup(context))
  141. BackupManager.Restore(context);
  142. }
  143. var nativePluginFolder = Path.Combine(context.DataPathDst, "Plugins");
  144. bool isFlat = Directory.Exists(nativePluginFolder) &&
  145. Directory.GetFiles(nativePluginFolder).Any(f => f.EndsWith(".dll"));
  146. bool force = !BackupManager.HasBackup(context) || ArgForce;
  147. var architecture = DetectArchitecture(context.Executable);
  148. Console.ForegroundColor = ConsoleColor.DarkCyan;
  149. Console.WriteLine("Installing files... ");
  150. CopyAll(new DirectoryInfo(context.DataPathSrc), new DirectoryInfo(context.DataPathDst), force,
  151. backup,
  152. (from, to) => NativePluginInterceptor(from, to, new DirectoryInfo(nativePluginFolder), isFlat,
  153. architecture));
  154. CopyAll(new DirectoryInfo(context.LibsPathSrc), new DirectoryInfo(context.LibsPathDst), force,
  155. backup,
  156. (from, to) => NativePluginInterceptor(from, to, new DirectoryInfo(nativePluginFolder), isFlat,
  157. architecture));
  158. CopyAll(new DirectoryInfo(context.IPARoot), new DirectoryInfo(context.ProjectRoot), force,
  159. backup,
  160. null, false);
  161. }
  162. else
  163. {
  164. Console.ForegroundColor = ConsoleColor.Yellow;
  165. Console.WriteLine("Not copying files because newer version already installed");
  166. }
  167. #region Create Plugin Folder
  168. if (!Directory.Exists(context.PluginsFolder))
  169. {
  170. Console.ForegroundColor = ConsoleColor.DarkYellow;
  171. Console.WriteLine("Creating plugins folder... ");
  172. Directory.CreateDirectory(context.PluginsFolder);
  173. Console.ResetColor();
  174. }
  175. #endregion
  176. }
  177. catch (Exception e)
  178. {
  179. Console.ForegroundColor = ConsoleColor.Red;
  180. Fail("Oops! This should not have happened.\n\n" + e);
  181. }
  182. Console.ResetColor();
  183. }
  184. private static void Revert(PatchContext context)
  185. {
  186. Console.ForegroundColor = ConsoleColor.Cyan;
  187. Console.Write("Restoring backup... ");
  188. if (BackupManager.Restore(context))
  189. {
  190. Console.WriteLine("Done!");
  191. }
  192. else
  193. {
  194. Console.WriteLine("Already vanilla or you removed your backups!");
  195. }
  196. if (File.Exists(context.ShortcutPath))
  197. {
  198. Console.WriteLine("Deleting shortcut...");
  199. File.Delete(context.ShortcutPath);
  200. }
  201. Console.WriteLine("");
  202. Console.WriteLine("--- Done reverting ---");
  203. Console.ResetColor();
  204. }
  205. private static void StartIfNeedBe(PatchContext context)
  206. {
  207. if (ArgStart.HasValue)
  208. {
  209. Process.Start(context.Executable, ArgStart.Value);
  210. }
  211. else
  212. {
  213. var argList = Arguments.CmdLine.PositionalArgs.ToList();
  214. argList.Remove(context.Executable);
  215. if (ArgLaunch)
  216. {
  217. Process.Start(context.Executable, Args(argList.ToArray()));
  218. }
  219. }
  220. }
  221. public static IEnumerable<FileInfo> NativePluginInterceptor(FileInfo from, FileInfo to,
  222. DirectoryInfo nativePluginFolder, bool isFlat, Architecture preferredArchitecture)
  223. {
  224. if (to.FullName.StartsWith(nativePluginFolder.FullName))
  225. {
  226. var relevantBit = to.FullName.Substring(nativePluginFolder.FullName.Length + 1);
  227. // Goes into the plugin folder!
  228. bool isFileFlat = !relevantBit.StartsWith("x86");
  229. if (isFlat && !isFileFlat)
  230. {
  231. // Flatten structure
  232. bool is64Bit = relevantBit.StartsWith("x86_64");
  233. if (!is64Bit && preferredArchitecture == Architecture.x86)
  234. {
  235. // 32 bit
  236. yield return new FileInfo(Path.Combine(nativePluginFolder.FullName,
  237. relevantBit.Substring("x86".Length + 1)));
  238. }
  239. else if (is64Bit && (preferredArchitecture == Architecture.x64 ||
  240. preferredArchitecture == Architecture.Unknown))
  241. {
  242. // 64 bit
  243. yield return new FileInfo(Path.Combine(nativePluginFolder.FullName,
  244. relevantBit.Substring("x86_64".Length + 1)));
  245. }
  246. }
  247. else if (!isFlat && isFileFlat)
  248. {
  249. // Deepen structure
  250. yield return new FileInfo(Path.Combine(Path.Combine(nativePluginFolder.FullName, "x86"),
  251. relevantBit));
  252. yield return new FileInfo(Path.Combine(Path.Combine(nativePluginFolder.FullName, "x86_64"),
  253. relevantBit));
  254. }
  255. else
  256. {
  257. yield return to;
  258. }
  259. }
  260. else
  261. {
  262. yield return to;
  263. }
  264. }
  265. public static void ClearLine()
  266. {
  267. if (IsConsole)
  268. {
  269. Console.SetCursorPosition(0, Console.CursorTop);
  270. int tpos = Console.CursorTop;
  271. Console.Write(new string(' ', Console.WindowWidth));
  272. Console.SetCursorPosition(0, tpos);
  273. }
  274. }
  275. private static IEnumerable<FileInfo> PassThroughInterceptor(FileInfo from, FileInfo to)
  276. {
  277. yield return to;
  278. }
  279. public static void CopyAll(DirectoryInfo source, DirectoryInfo target, bool aggressive, BackupUnit backup,
  280. Func<FileInfo, FileInfo, IEnumerable<FileInfo>> interceptor = null, bool recurse = true)
  281. {
  282. if (interceptor == null)
  283. {
  284. interceptor = PassThroughInterceptor;
  285. }
  286. // Copy each file into the new directory.
  287. foreach (var fi in source.GetFiles())
  288. {
  289. foreach (var targetFile in interceptor(fi, new FileInfo(Path.Combine(target.FullName, fi.Name))))
  290. {
  291. if (targetFile.Exists && targetFile.LastWriteTimeUtc >= fi.LastWriteTimeUtc && !aggressive)
  292. continue;
  293. Debug.Assert(targetFile.Directory != null, "targetFile.Directory != null");
  294. targetFile.Directory?.Create();
  295. LineBack();
  296. ClearLine();
  297. Console.WriteLine(@"Copying {0}", targetFile.FullName);
  298. backup.Add(targetFile);
  299. fi.CopyTo(targetFile.FullName, true);
  300. }
  301. }
  302. // Copy each subdirectory using recursion.
  303. if (!recurse) return;
  304. foreach (var diSourceSubDir in source.GetDirectories())
  305. {
  306. var nextTargetSubDir = new DirectoryInfo(Path.Combine(target.FullName, diSourceSubDir.Name));
  307. CopyAll(diSourceSubDir, nextTargetSubDir, aggressive, backup, interceptor);
  308. }
  309. }
  310. private static void Fail(string message)
  311. {
  312. Console.Error.WriteLine("ERROR: " + message);
  313. WaitForEnd();
  314. Environment.Exit(1);
  315. }
  316. public static string Args(params string[] args)
  317. {
  318. return string.Join(" ", args.Select(EncodeParameterArgument).ToArray());
  319. }
  320. /// <summary>
  321. /// Encodes an argument for passing into a program
  322. /// </summary>
  323. /// <param name="original">The value_ that should be received by the program</param>
  324. /// <returns>The value_ which needs to be passed to the program for the original value_
  325. /// to come through</returns>
  326. public static string EncodeParameterArgument(string original)
  327. {
  328. if (string.IsNullOrEmpty(original))
  329. return original;
  330. string value = Regex.Replace(original, @"(\\*)" + "\"", @"$1\$0");
  331. value = Regex.Replace(value, @"^(.*\s.*?)(\\*)$", "\"$1$2$2\"");
  332. return value;
  333. }
  334. public static Architecture DetectArchitecture(string assembly)
  335. {
  336. using (var reader = new BinaryReader(File.OpenRead(assembly)))
  337. {
  338. var header = reader.ReadUInt16();
  339. if (header == 0x5a4d)
  340. {
  341. reader.BaseStream.Seek(60, SeekOrigin.Begin); // this location contains the offset for the PE header
  342. var peOffset = reader.ReadUInt32();
  343. reader.BaseStream.Seek(peOffset + 4, SeekOrigin.Begin);
  344. var machine = reader.ReadUInt16();
  345. if (machine == 0x8664) // IMAGE_FILE_MACHINE_AMD64
  346. return Architecture.x64;
  347. if (machine == 0x014c) // IMAGE_FILE_MACHINE_I386
  348. return Architecture.x86;
  349. if (machine == 0x0200) // IMAGE_FILE_MACHINE_IA64
  350. return Architecture.x64;
  351. return Architecture.Unknown;
  352. }
  353. // Not a supported binary
  354. return Architecture.Unknown;
  355. }
  356. }
  357. public static void ResetLine()
  358. {
  359. if (IsConsole)
  360. Console.CursorLeft = 0;
  361. else
  362. Console.Write("\r");
  363. }
  364. public static void LineBack()
  365. {
  366. if (IsConsole)
  367. Console.CursorTop--;
  368. else
  369. Console.Write("\x1b[1A");
  370. }
  371. [DllImport("kernel32.dll")]
  372. private static extern IntPtr GetConsoleWindow();
  373. private static bool? isConsole;
  374. public static bool IsConsole
  375. {
  376. get
  377. {
  378. if (isConsole == null)
  379. isConsole = GetConsoleWindow() != IntPtr.Zero;
  380. return isConsole.Value;
  381. }
  382. }
  383. public abstract class Keyboard
  384. {
  385. [Flags]
  386. private enum KeyStates
  387. {
  388. None = 0,
  389. Down = 1,
  390. Toggled = 2
  391. }
  392. [DllImport("user32.dll", CharSet = CharSet.Auto, ExactSpelling = true)]
  393. private static extern short GetKeyState(int keyCode);
  394. private static KeyStates KeyState(Keys key)
  395. {
  396. KeyStates state = KeyStates.None;
  397. short retVal = GetKeyState((int)key);
  398. //If the high-order bit is 1, the key is down
  399. //otherwise, it is up.
  400. if ((retVal & 0x8000) == 0x8000)
  401. state |= KeyStates.Down;
  402. //If the low-order bit is 1, the key is toggled.
  403. if ((retVal & 1) == 1)
  404. state |= KeyStates.Toggled;
  405. return state;
  406. }
  407. public static bool IsKeyDown(Keys key)
  408. {
  409. return KeyStates.Down == (KeyState(key) & KeyStates.Down);
  410. }
  411. public static bool IsKeyToggled(Keys key)
  412. {
  413. return KeyStates.Toggled == (KeyState(key) & KeyStates.Toggled);
  414. }
  415. }
  416. }
  417. }