Skip to content

Commit 80f982e

Browse files
raonygamerraonygamer
authored
Merge pull request #5 from AmethystAPI/new_module_tweaker
New module tweaker
2 parents 11663f9 + bcd892c commit 80f982e

38 files changed

Lines changed: 933 additions & 566 deletions

‎Common/Diagnostics/Logger.cs‎

Lines changed: 13 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
using System.Reflection;
1+
using Amethyst.Common.Utility;
2+
using System.Reflection;
23

34
namespace Amethyst.Common.Diagnostics
45
{
@@ -10,43 +11,28 @@ public static void WriteLine(object? message = null, ConsoleColor color = Consol
1011
Console.WriteLine(message);
1112
Console.ResetColor();
1213
}
13-
14-
public static void Info(object? message)
14+
15+
public static void Info(object? message, CursorLocation? location = null)
1516
{
16-
WriteLine($"{(
17-
Assembly.GetEntryAssembly()?.GetName()?.Name is { } name ?
18-
$"[{name}] " :
19-
"")}[INFO] {message}", ConsoleColor.White);
17+
WriteLine($"{(location is not null && location.File != "<unknown>" ? location.ToString() + ": " : (Assembly.GetEntryAssembly()?.GetName().Name is not string name ? "Unknown: " : name.Trim() + ": "))} message: {message}", ConsoleColor.White);
2018
}
2119

22-
public static void Debug(object? message)
20+
public static void Debug(object? message, CursorLocation? location = null)
2321
{
2422
#if DEBUG
25-
WriteLine($"{(
26-
Assembly.GetEntryAssembly()?.GetName()?.Name is { } name ?
27-
$"[{name}] " :
28-
"")}[DEBUG] {message}", ConsoleColor.White);
23+
WriteLine($"{(location is not null && location.File != "<unknown>" ? location.ToString() + ": " : (Assembly.GetEntryAssembly()?.GetName().Name is not string name ? "Unknown: " : name.Trim() + ": "))}message: {message}", ConsoleColor.White);
2924
#endif
3025
}
3126

32-
public static void Warn(string message) =>
33-
WriteLine($"{(
34-
Assembly.GetEntryAssembly()?.GetName()?.Name is { } name ?
35-
$"[{name}] " :
36-
"")}[WARN] {message}", ConsoleColor.Yellow);
27+
public static void Warn(string message, CursorLocation? location = null) =>
28+
WriteLine($"{(location is not null && location.File != "<unknown>" ? location?.ToString() + ": " : (Assembly.GetEntryAssembly()?.GetName().Name is not string name ? "Unknown: " : name.Trim() + ": "))}warning: {message}", ConsoleColor.Yellow);
3729

38-
public static void Error(string message) =>
39-
WriteLine($"{(
40-
Assembly.GetEntryAssembly()?.GetName()?.Name is { } name ?
41-
$"[{name}] " :
42-
"")}[ERROR] {message}", ConsoleColor.Red);
30+
public static void Error(string message, CursorLocation? location = null) =>
31+
WriteLine($"{(location is not null && location.File != "<unknown>" ? location?.ToString() + ": " : (Assembly.GetEntryAssembly()?.GetName().Name is not string name ? "Unknown: " : name.Trim() + ": "))}error: {message}", ConsoleColor.Red);
4332

44-
public static void Fatal(string message)
33+
public static void Fatal(string message, CursorLocation? location = null)
4534
{
46-
WriteLine($"{(
47-
Assembly.GetEntryAssembly()?.GetName()?.Name is { } name ?
48-
$"[{name}] " :
49-
"")}[FATAL] {message}", ConsoleColor.Magenta);
35+
WriteLine($"{(location is not null && location.File != "<unknown>" ? location?.ToString() + ": " : (Assembly.GetEntryAssembly()?.GetName().Name is not string name ? "Unknown: " : name.Trim() + ": "))}fatal error: {message}", ConsoleColor.Magenta);
5036
Environment.Exit(1);
5137
}
5238
}

‎Common/Models/VariableSymbolModel.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,8 @@ public class VariableSymbolModel
1414

1515
[JsonProperty("address")]
1616
public string Address { get; set; } = string.Empty;
17+
18+
[JsonProperty("is_vaddress")]
19+
public bool IsVirtualTableAddress { get; set; } = false;
1720
}
1821
}

‎Common/Models/VirtualFunctionSymbolModel.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,5 +23,8 @@ public class VirtualFunctionSymbolModel
2323

2424
[JsonProperty("index")]
2525
public uint Index { get; set; } = 0;
26+
27+
[JsonProperty("is_vdtor")]
28+
public bool IsVirtualDestructor { get; set; } = false;
2629
}
2730
}

‎Common/Tracking/FileTracker.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ public FileTracker(DirectoryInfo inputDirectory, FileInfo checksumFile, string[]
7272
if (Filters.Any() && !Filters.Any(f => Path.GetRelativePath(InputDirectory.FullName, file.FullName).StartsWith(f)))
7373
continue;
7474
string filePath = file.FullName.NormalizeSlashes();
75-
#if !DEBUG
75+
#if DEBUG
7676
string content = File.ReadAllText(file.FullName);
7777
ulong hash = XXH64.DigestOf(Encoding.UTF8.GetBytes(content));
7878

‎Common/Utility/CursorLocation.cs‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
using Amethyst.Common.Extensions;
2+
3+
namespace Amethyst.Common.Utility {
4+
public class CursorLocation
5+
{
6+
public string File { get; set; }
7+
public uint Line { get; set; }
8+
public uint Column { get; set; }
9+
10+
public CursorLocation(string file, uint line, uint column)
11+
{
12+
if (string.IsNullOrEmpty(file) || !System.IO.File.Exists(file))
13+
File = "<unknown>";
14+
else
15+
File = Path.GetFullPath(file).NormalizeSlashes();
16+
Line = line;
17+
Column = column;
18+
}
19+
20+
override public string ToString()
21+
{
22+
if (File == "<unknown>")
23+
return "";
24+
return $"{File}({Line},{Column})";
25+
}
26+
}
27+
}

‎Common/Utility/Utils.cs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,5 +59,24 @@ public static void CreateDefinitionFile(string defFile, IEnumerable<string> mang
5959
sb.AppendLine("; End of generated file.");
6060
File.WriteAllText(defFile, sb.ToString());
6161
}
62+
63+
public static void WritePrefixedString(this BinaryWriter writer, string str) {
64+
byte[] bytes = Encoding.UTF8.GetBytes(str);
65+
writer.Write(bytes.Length);
66+
writer.Write(bytes);
67+
}
68+
69+
public static string ReadPrefixedString(this BinaryReader reader) {
70+
int length = reader.ReadInt32();
71+
byte[] bytes = reader.ReadBytes(length);
72+
return Encoding.UTF8.GetString(bytes);
73+
}
74+
75+
public static void Align(this BinaryWriter writer, int alignment = 16, byte pad = 0x00) {
76+
long pos = writer.BaseStream.Position;
77+
int padding = (int)((alignment - (pos % alignment)) % alignment);
78+
for (int i = 0; i < padding; i++)
79+
writer.Write(pad);
80+
}
6281
}
6382
}

‎ModuleTweaker/Amethyst.ModuleTweaker.csproj‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,17 +7,17 @@
77
<Nullable>enable</Nullable>
88
<AssemblyName>Amethyst.ModuleTweaker</AssemblyName>
99
<RootNamespace>Amethyst.ModuleTweaker</RootNamespace>
10-
<Version>1.0.6</Version>
10+
<Version>2.0.0</Version>
1111
</PropertyGroup>
1212

13-
<ItemGroup>
14-
<ProjectReference Include="..\Common\Amethyst.Common.csproj" />
15-
</ItemGroup>
16-
1713
<ItemGroup>
1814
<PackageReference Include="AsmResolver" Version="5.5.1" />
1915
<PackageReference Include="AsmResolver.PE" Version="5.5.1" />
2016
<PackageReference Include="AsmResolver.PE.File" Version="5.5.1" />
2117
</ItemGroup>
2218

19+
<ItemGroup>
20+
<ProjectReference Include="..\Common\Amethyst.Common.csproj" />
21+
</ItemGroup>
22+
2323
</Project>

‎ModuleTweaker/Commands/MainCommand.cs‎

Lines changed: 97 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2,44 +2,65 @@
22
using Amethyst.Common.Models;
33
using Amethyst.ModuleTweaker.Patching;
44
using AsmResolver.PE.File;
5+
using AsmResolver.PE.Imports;
56
using CliFx;
67
using CliFx.Attributes;
78
using CliFx.Infrastructure;
9+
using K4os.Hash.xxHash;
810
using Newtonsoft.Json;
11+
using System.Globalization;
912

1013
namespace Amethyst.ModuleTweaker.Commands
1114
{
1215
[Command(Description = "Patches or unpatches modules for runtime importing support.")]
1316
public class MainCommand : ICommand
1417
{
15-
[CommandOption("module", 'm', Description = "The specified module path to patch.")]
18+
[CommandOption("module", 'm', Description = "The specified module path to patch.", IsRequired = true)]
1619
public string ModulePath { get; set; } = null!;
1720

18-
[CommandOption("symbols", 's', Description = "Path to directory containing *.symbols.json to use for patching.")]
21+
[CommandOption("symbols", 's', Description = "Path to directory containing *.symbols.json to use for patching.", IsRequired = true)]
1922
public string SymbolsPath { get; set; } = null!;
2023

24+
[CommandOption("output", 'o', Description = "Path to save temporary files, don't confuse with -m.")]
25+
public string OutputPath { get; set; } = null!;
26+
2127
public ValueTask ExecuteAsync(IConsole console)
2228
{
2329
FileInfo module = new(ModulePath);
24-
DirectoryInfo symbols = new(SymbolsPath);
25-
if (module.Exists is false)
26-
{
27-
Logger.Warn("Couldn't patch module, specified module does not exist.");
30+
DirectoryInfo symbolsDir = new(SymbolsPath);
31+
if (module.Exists is false) {
32+
Logger.Fatal("Couldn't patch module, specified module does not exist.");
2833
return default;
2934
}
3035

31-
if (symbols.Exists is false)
32-
{
33-
Logger.Warn("Couldn't patch module, specified symbols directory does not exist.");
36+
if (symbolsDir.Exists is false) {
37+
Logger.Fatal("Couldn't patch module, specified symbols directory does not exist.");
3438
return default;
3539
}
3640

41+
if (string.IsNullOrEmpty(OutputPath)) {
42+
OutputPath = Path.GetFullPath(Path.Combine(SymbolsPath, "../"));
43+
}
44+
DirectoryInfo outDir = new(OutputPath);
45+
46+
ulong ParseAddress(string? address)
47+
{
48+
if (string.IsNullOrEmpty(address))
49+
return 0x0;
50+
if (address.StartsWith("0x", StringComparison.OrdinalIgnoreCase))
51+
address = address[2..];
52+
if (!ulong.TryParse(address, NumberStyles.HexNumber, null, out var addr))
53+
return 0x0;
54+
return addr;
55+
}
56+
57+
SymbolFactory.Register(new SymbolType(1, "pe32+", "data"), () => new Patching.PE.V1.PEDataSymbol());
58+
SymbolFactory.Register(new SymbolType(1, "pe32+", "function"), () => new Patching.PE.V1.PEFunctionSymbol());
59+
HeaderFactory.Register(new HeaderType(1, "pe32+"), (args) => new Patching.PE.V1.PEImporterHeader());
60+
3761
// Collect all symbol files and accumulate mangled names
38-
IEnumerable<FileInfo> symbolFiles = symbols.EnumerateFiles("*.json", SearchOption.AllDirectories);
39-
HashSet<FunctionSymbolModel> methods = [];
40-
HashSet<VariableSymbolModel> variables = [];
41-
HashSet<VirtualTableSymbolModel> vtables = [];
42-
HashSet<VirtualFunctionSymbolModel> vfuncs = [];
62+
IEnumerable<FileInfo> symbolFiles = symbolsDir.EnumerateFiles("*.json", SearchOption.AllDirectories);
63+
List<AbstractSymbol> symbols = [];
4364
foreach (var symbolFile in symbolFiles)
4465
{
4566
using var stream = symbolFile.OpenRead();
@@ -50,29 +71,48 @@ public ValueTask ExecuteAsync(IConsole console)
5071
switch (symbolJson.FormatVersion)
5172
{
5273
case 1:
53-
foreach (var function in symbolJson.Functions)
54-
{
74+
foreach (var function in symbolJson.Functions) {
5575
if (string.IsNullOrEmpty(function.Name))
5676
continue;
57-
methods.Add(function);
77+
symbols.Add(new Patching.PE.V1.PEFunctionSymbol {
78+
Name = function.Name,
79+
IsVirtual = false,
80+
IsSignature = function.Signature is not null,
81+
Address = ParseAddress(function.Address),
82+
Signature = function.Signature ?? string.Empty
83+
});
5884
}
59-
foreach (var variable in symbolJson.Variables)
60-
{
61-
if (string.IsNullOrEmpty(variable.Name))
85+
foreach (var vfunc in symbolJson.VirtualFunctions) {
86+
if (string.IsNullOrEmpty(vfunc.Name))
6287
continue;
63-
variables.Add(variable);
88+
symbols.Add(new Patching.PE.V1.PEFunctionSymbol {
89+
Name = vfunc.Name,
90+
IsVirtual = true,
91+
VirtualIndex = vfunc.Index,
92+
VirtualTable = vfunc.VirtualTable ?? "this",
93+
IsDestructor = vfunc.IsVirtualDestructor,
94+
HasStorage = vfunc.IsVirtualDestructor
95+
});
6496
}
65-
foreach (var vtable in symbolJson.VirtualTables)
66-
{
67-
if (string.IsNullOrEmpty(vtable.Name))
97+
foreach (var variable in symbolJson.Variables) {
98+
if (string.IsNullOrEmpty(variable.Name))
6899
continue;
69-
vtables.Add(vtable);
100+
symbols.Add(new Patching.PE.V1.PEDataSymbol {
101+
Name = variable.Name,
102+
IsVirtualTable = false,
103+
Address = ParseAddress(variable.Address),
104+
IsVirtualTableAddress = variable.IsVirtualTableAddress,
105+
HasStorage = variable.IsVirtualTableAddress
106+
});
70107
}
71-
foreach (var vfunc in symbolJson.VirtualFunctions)
72-
{
73-
if (string.IsNullOrEmpty(vfunc.Name))
108+
foreach (var vtable in symbolJson.VirtualTables) {
109+
if (string.IsNullOrEmpty(vtable.Name))
74110
continue;
75-
vfuncs.Add(vfunc);
111+
symbols.Add(new Patching.PE.V1.PEDataSymbol {
112+
Name = vtable.Name,
113+
IsVirtualTable = true,
114+
Address = ParseAddress(vtable.Address)
115+
});
76116
}
77117
break;
78118
}
@@ -82,13 +122,35 @@ public ValueTask ExecuteAsync(IConsole console)
82122
try
83123
{
84124
// Patch the module
85-
var file = PEFile.FromFile(ModulePath);
86-
PEFileHelper helper = new(file);
87-
if (helper.Patch(methods, variables, vtables, vfuncs))
125+
var bytes = File.ReadAllBytes(ModulePath);
126+
ulong hash = XXH64.DigestOf(bytes);
127+
if (File.Exists(Path.Combine(outDir.FullName, "module_hash.txt"))) {
128+
var existingHash = File.ReadAllText(Path.Combine(outDir.FullName, "module_hash.txt"));
129+
if (ulong.TryParse(existingHash, NumberStyles.HexNumber, null, out var existingHashValue)) {
130+
if (existingHashValue == hash) {
131+
Logger.Info("Module hash matches previous hash, skipping patch.");
132+
return default;
133+
}
134+
}
135+
}
136+
137+
var peFile = PEFile.FromBytes(bytes);
138+
if (peFile is null) {
139+
Logger.Fatal("Failed to read module as a PE file.");
140+
return default;
141+
}
142+
Logger.Info($"Loaded module '{ModulePath}' as PE file.");
143+
var patcher = new Patching.PE.PEPatcher(peFile, symbols);
144+
145+
if (patcher.Patch())
88146
{
89-
file.AlignSections();
90-
File.Copy(ModulePath, ModulePath + ".backup", true);
91-
file.Write(ModulePath);
147+
File.Copy(ModulePath, ModulePath + ".bak", true);
148+
using var ms = new MemoryStream();
149+
peFile.Write(ms);
150+
var newBytes = ms.ToArray();
151+
ulong newHash = XXH64.DigestOf(newBytes);
152+
File.WriteAllBytes(ModulePath, newBytes);
153+
File.WriteAllText(Path.Combine(outDir.FullName, "module_hash.txt"), newHash.ToString("X16"));
92154
}
93155
}
94156
catch (Exception ex)

0 commit comments

Comments
 (0)