LenovoLegionToolkit

1
using System;
2
using System.Runtime.InteropServices;
3
using System.Threading.Tasks;
4
using LenovoLegionToolkit.Lib.Utils;
5
using Windows.Win32;
6
using Windows.Win32.Security;
7

8
namespace LenovoLegionToolkit.Lib.Features;
9

10
public abstract class AbstractUEFIFeature<T> : IFeature<T> where T : struct, Enum, IComparable
11
{
12
    private readonly string _guid;
13
    private readonly string _scopeName;
14
    private readonly uint _scopeAttribute;
15

16
    protected AbstractUEFIFeature(string guid, string scopeName, uint scopeAttribute)
17
    {
18
        _guid = guid;
19
        _scopeName = scopeName;
20
        _scopeAttribute = scopeAttribute;
21
    }
22

23
    public async Task<bool> IsSupportedAsync()
24
    {
25
        try
26
        {
27
            _ = await GetStateAsync().ConfigureAwait(false);
28
            return true;
29
        }
30
        catch
31
        {
32
            return false;
33
        }
34
    }
35

36
    public Task<T[]> GetAllStatesAsync() => Task.FromResult(Enum.GetValues<T>());
37

38
    public abstract Task<T> GetStateAsync();
39

40
    public abstract Task SetStateAsync(T state);
41

42
    protected unsafe Task<TS> ReadFromUefiAsync<TS>() where TS : struct => Task.Run(() =>
43
    {
44
        if (Log.Instance.IsTraceEnabled)
45
            Log.Instance.Trace($"Reading from UEFI... [feature={GetType().Name}]");
46

47
        var ptr = Marshal.AllocHGlobal(Marshal.SizeOf<TS>());
48

49
        try
50
        {
51
            if (!SetPrivilege(true))
52
            {
53
                if (Log.Instance.IsTraceEnabled)
54
                    Log.Instance.Trace($"Cannot set UEFI privileges [feature={GetType().Name}]");
55

56
                throw new InvalidOperationException("Cannot set privileges UEFI");
57
            }
58

59
            var ptrSize = (uint)Marshal.SizeOf<TS>();
60
            if (PInvoke.GetFirmwareEnvironmentVariableEx(_scopeName, _guid, ptr.ToPointer(), ptrSize, null) != 0)
61
            {
62
                var result = Marshal.PtrToStructure<TS>(ptr);
63

64
                if (Log.Instance.IsTraceEnabled)
65
                    Log.Instance.Trace($"Read from UEFI successful [feature={GetType().Name}]");
66

67
                return result;
68
            }
69
            else
70
            {
71
                if (Log.Instance.IsTraceEnabled)
72
                    Log.Instance.Trace($"Cannot read variable {_scopeName} from UEFI [feature={GetType().Name}]");
73

74
                throw new InvalidOperationException($"Cannot read variable {_scopeName} from UEFI");
75
            }
76
        }
77
        finally
78
        {
79
            Marshal.FreeHGlobal(ptr);
80
            SetPrivilege(false);
81
        }
82
    });
83

84
    protected unsafe Task WriteToUefiAsync<TS>(TS structure) where TS : struct => Task.Run(() =>
85
    {
86
        var ptr = Marshal.AllocHGlobal(Marshal.SizeOf<TS>());
87

88
        try
89
        {
90
            if (!SetPrivilege(true))
91
            {
92
                if (Log.Instance.IsTraceEnabled)
93
                    Log.Instance.Trace($"Cannot set UEFI privileges [feature={GetType().Name}]");
94

95
                throw new InvalidOperationException("Cannot set UEFI privileges");
96
            }
97

98
            Marshal.StructureToPtr(structure, ptr, false);
99
            var ptrSize = (uint)Marshal.SizeOf<TS>();
100
            if (!PInvoke.SetFirmwareEnvironmentVariableEx(_scopeName, _guid, ptr.ToPointer(), ptrSize, _scopeAttribute))
101
            {
102
                if (Log.Instance.IsTraceEnabled)
103
                    Log.Instance.Trace($"Cannot write variable {_scopeName} to UEFI [feature={GetType().Name}]");
104

105
                throw new InvalidOperationException($"Cannot write variable {_scopeName} to UEFI");
106
            }
107
            else
108
            {
109
                if (Log.Instance.IsTraceEnabled)
110
                    Log.Instance.Trace($"WriteAsync to UEFI successful [feature={GetType().Name}]");
111
            }
112
        }
113
        finally
114
        {
115
            Marshal.FreeHGlobal(ptr);
116
            SetPrivilege(false);
117
        }
118
    });
119

120
    private unsafe bool SetPrivilege(bool enable)
121
    {
122
        try
123
        {
124
            using var handle = PInvoke.GetCurrentProcess_SafeHandle();
125

126
            if (!PInvoke.OpenProcessToken(handle, TOKEN_ACCESS_MASK.TOKEN_QUERY | TOKEN_ACCESS_MASK.TOKEN_ADJUST_PRIVILEGES, out var token))
127
            {
128
                if (Log.Instance.IsTraceEnabled)
129
                    Log.Instance.Trace($"Could not open process token [feature={GetType().Name}]");
130

131
                return false;
132
            }
133

134
            if (!PInvoke.LookupPrivilegeValue(null, "SeSystemEnvironmentPrivilege", out var luid))
135
            {
136
                if (Log.Instance.IsTraceEnabled)
137
                    Log.Instance.Trace($"Could not look up privilege value [feature={GetType().Name}]");
138

139
                return false;
140
            }
141

142
            var state = new TOKEN_PRIVILEGES { PrivilegeCount = 1 };
143
            state.Privileges._0 = new LUID_AND_ATTRIBUTES
144
            {
145
                Luid = luid,
146
                Attributes = enable ? TOKEN_PRIVILEGES_ATTRIBUTES.SE_PRIVILEGE_ENABLED : 0
147
            };
148

149
            if (!PInvoke.AdjustTokenPrivileges(token, false, state, 0, null, null))
150
            {
151
                if (Log.Instance.IsTraceEnabled)
152
                    Log.Instance.Trace($"Could not adjust token privileges [feature={GetType().Name}]");
153

154
                return false;
155
            }
156

157
            return true;
158
        }
159
        catch (Exception ex)
160
        {
161
            if (Log.Instance.IsTraceEnabled)
162
                Log.Instance.Trace($"Exception while setting privilege. [feature={GetType().Name}]", ex);
163

164
            return false;
165
        }
166
    }
167
}
168