···
1
+
#!@python3@/bin/python3 -B
3
+
from dataclasses import dataclass
4
+
from typing import Any, Callable, Dict, List, Optional, Tuple
8
+
from ctypes import CDLL
18
+
libc = CDLL("libc.so.6")
19
+
install_config = json.load(open('@configPath@', 'r'))
22
+
def config(*path: str) -> Optional[Any]:
23
+
result = install_config
24
+
for component in path:
25
+
result = result[component]
29
+
def get_system_path(profile: str = 'system', gen: Optional[str] = None, spec: Optional[str] = None) -> str:
30
+
basename = f'{profile}-{gen}-link' if gen is not None else profile
31
+
profiles_dir = '/nix/var/nix/profiles'
32
+
if profile == 'system':
33
+
result = os.path.join(profiles_dir, basename)
35
+
result = os.path.join(profiles_dir, 'system-profiles', basename)
37
+
if spec is not None:
38
+
result = os.path.join(result, 'specialisation', spec)
43
+
def get_profiles() -> List[str]:
44
+
profiles_dir = '/nix/var/nix/profiles/system-profiles/'
45
+
dirs = os.listdir(profiles_dir) if os.path.isdir(profiles_dir) else []
47
+
return [path for path in dirs if not path.endswith('-link')]
50
+
def get_gens(profile: str = 'system') -> List[Tuple[int, List[str]]]:
51
+
nix_env = os.path.join(config('nixPath'), 'bin', 'nix-env')
52
+
output = subprocess.check_output([
53
+
nix_env, '--list-generations',
54
+
'-p', get_system_path(profile),
55
+
'--option', 'build-users-group', '',
56
+
], universal_newlines=True)
58
+
gen_lines = output.splitlines()
59
+
gen_nums = [int(line.split()[0]) for line in gen_lines]
61
+
return [gen for gen in gen_nums][-config('maxGenerations'):]
64
+
def is_encrypted(device: str) -> bool:
65
+
for name, _ in config('luksDevices'):
66
+
if os.readlink(os.path.join('/dev/mapper', name)) == os.readlink(device):
72
+
def is_fs_type_supported(fs_type: str) -> bool:
73
+
return fs_type.startswith('vfat')
78
+
def get_copied_path_uri(path: str, target: str) -> str:
79
+
package_id = os.path.basename(os.path.dirname(path))
80
+
suffix = os.path.basename(path)
81
+
dest_file = f'{package_id}-{suffix}'
82
+
dest_path = os.path.join(refind_dir, target, dest_file)
84
+
if not os.path.exists(dest_path):
85
+
copy_file(path, dest_path)
87
+
paths[dest_path] = True
89
+
return os.path.join('/efi/refind', target, dest_file)
91
+
def get_path_uri(path: str) -> str:
92
+
return get_copied_path_uri(path, "")
95
+
def get_file_uri(profile: str, gen: Optional[str], spec: Optional[str], name: str) -> str:
96
+
gen_path = get_system_path(profile, gen, spec)
97
+
path_in_store = os.path.realpath(os.path.join(gen_path, name))
98
+
return get_path_uri(path_in_store)
101
+
def get_kernel_uri(kernel_path: str) -> str:
102
+
return get_copied_path_uri(kernel_path, "kernels")
110
+
kernelParams: List[str]
113
+
specialisations: Dict[str, "BootSpec"]
114
+
initrd: str | None = None
115
+
initrdSecrets: str | None = None
118
+
def bootjson_to_bootspec(bootjson: dict) -> BootSpec:
119
+
specialisations = bootjson['org.nixos.specialisation.v1']
120
+
specialisations = {k: bootjson_to_bootspec(v) for k, v in specialisations.items()}
122
+
**bootjson['org.nixos.bootspec.v1'],
123
+
specialisations=specialisations,
127
+
def config_entry(is_sub: bool, bootspec: BootSpec, label: str, time: str) -> str:
132
+
entry += f'menuentry "{label}" {{\n'
133
+
entry += ' loader ' + get_kernel_uri(bootspec.kernel) + '\n'
135
+
if bootspec.initrd:
136
+
entry += ' initrd ' + get_kernel_uri(bootspec.initrd) + '\n'
138
+
entry += ' options "' + ' '.join(['init=' + bootspec.init] + bootspec.kernelParams).strip() + '"\n'
143
+
def generate_config_entry(profile: str, gen: str, special: bool, group_name: str) -> str:
144
+
time = datetime.datetime.fromtimestamp(os.stat(get_system_path(profile,gen), follow_symlinks=False).st_mtime).strftime("%F %H:%M:%S")
145
+
boot_json = json.load(open(os.path.join(get_system_path(profile, gen), 'boot.json'), 'r'))
146
+
boot_spec = bootjson_to_bootspec(boot_json)
148
+
specialisation_list = boot_spec.specialisations.items()
151
+
if len(specialisation_list) > 0:
152
+
entry += f'menuentry "NixOS {group_name} Generation {gen}" {{\n'
153
+
entry += config_entry(True, boot_spec, f'Default', str(time))
155
+
for spec, spec_boot_spec in specialisation_list:
156
+
entry += config_entry(True, spec_boot_spec, f'{spec}', str(time))
160
+
entry += config_entry(False, boot_spec, f'NixOS {group_name} Generation {gen}', str(time))
164
+
def find_disk_device(part: str) -> str:
165
+
part = os.path.realpath(part)
166
+
part = part.removeprefix('/dev/')
167
+
disk = os.path.realpath(os.path.join('/sys', 'class', 'block', part))
168
+
disk = os.path.dirname(disk)
170
+
return os.path.join('/dev', os.path.basename(disk))
173
+
def find_mounted_device(path: str) -> str:
174
+
path = os.path.abspath(path)
176
+
while not os.path.ismount(path):
177
+
path = os.path.dirname(path)
179
+
devices = [x for x in psutil.disk_partitions() if x.mountpoint == path]
181
+
assert len(devices) == 1
182
+
return devices[0].device
185
+
def copy_file(from_path: str, to_path: str):
186
+
dirname = os.path.dirname(to_path)
188
+
if not os.path.exists(dirname):
189
+
os.makedirs(dirname)
191
+
shutil.copyfile(from_path, to_path + ".tmp")
192
+
os.rename(to_path + ".tmp", to_path)
194
+
paths[to_path] = True
197
+
def install_bootloader() -> None:
200
+
refind_dir = os.path.join(str(config('efiMountPoint')), 'efi', 'refind')
202
+
if not os.path.exists(refind_dir):
203
+
os.makedirs(refind_dir)
205
+
for dir, dirs, files in os.walk(refind_dir, topdown=True):
207
+
paths[os.path.join(dir, file)] = False
209
+
profiles = [('system', get_gens())]
211
+
for profile in get_profiles():
212
+
profiles += [(profile, get_gens(profile))]
214
+
timeout = config('timeout')
216
+
last_gen = get_gens()[-1]
217
+
last_gen_json = json.load(open(os.path.join(get_system_path('system', last_gen), 'boot.json'), 'r'))
218
+
last_gen_boot_spec = bootjson_to_bootspec(last_gen_json)
220
+
config_file = str(config('extraConfig')) + '\n'
221
+
config_file += textwrap.dedent(f'''
223
+
default_selection {3 if len(last_gen_boot_spec.specialisations.items()) > 0 else 2}
226
+
config_file += textwrap.dedent('''
227
+
# NixOS boot entries start here
230
+
for (profile, gens) in profiles:
231
+
group_name = 'default profile' if profile == 'system' else f"profile '{profile}'"
234
+
for gen in sorted(gens, key=lambda x: x, reverse=True):
235
+
config_file += generate_config_entry(profile, gen, isFirst, group_name)
238
+
config_file_path = os.path.join(refind_dir, 'refind.conf')
239
+
config_file += '\n# NixOS boot entries end here\n\n'
241
+
with open(f"{config_file_path}.tmp", 'w') as file:
243
+
file.write(config_file.strip())
245
+
os.fsync(file.fileno())
246
+
os.rename(f"{config_file_path}.tmp", config_file_path)
248
+
paths[config_file_path] = True
250
+
for dest_path, source_path in config('additionalFiles').items():
251
+
dest_path = os.path.join(refind_dir, dest_path)
253
+
copy_file(source_path, dest_path)
255
+
cpu_family = config('hostArchitecture', 'family')
256
+
if cpu_family == 'x86':
257
+
if config('hostArchitecture', 'bits') == 32:
258
+
boot_file = 'BOOTIA32.EFI'
259
+
efi_file = 'refind_ia32.efi'
260
+
elif config('hostArchitecture', 'bits') == 64:
261
+
boot_file = 'BOOTX64.EFI'
262
+
efi_file = 'refind_x64.efi'
263
+
elif cpu_family == 'arm':
264
+
if config('hostArchitecture', 'arch') == 'armv8-a' and config('hostArchitecture', 'bits') == 64:
265
+
boot_file = 'BOOTAA64.EFI'
266
+
efi_file = 'refind_aa64.efi'
268
+
raise Exception(f'Unsupported CPU arch: {config("hostArchitecture", "arch")}')
270
+
raise Exception(f'Unsupported CPU family: {cpu_family}')
272
+
efi_path = os.path.join(config('refindPath'), 'share', 'refind', efi_file)
273
+
dest_path = os.path.join(config('efiMountPoint'), 'efi', 'boot' if config('efiRemovable') else 'refind', boot_file)
275
+
copy_file(efi_path, dest_path)
277
+
if not config('efiRemovable') and not config('canTouchEfiVariables'):
278
+
print('warning: boot.loader.efi.canTouchEfiVariables is set to false while boot.loader.limine.efiInstallAsRemovable.\n This may render the system unbootable.')
280
+
if config('canTouchEfiVariables'):
281
+
if config('efiRemovable'):
282
+
print('note: boot.loader.limine.efiInstallAsRemovable is true, no need to add EFI entry.')
284
+
efibootmgr = os.path.join(str(config('efiBootMgrPath')), 'bin', 'efibootmgr')
285
+
efi_partition = find_mounted_device(str(config('efiMountPoint')))
286
+
efi_disk = find_disk_device(efi_partition)
288
+
efibootmgr_output = subprocess.check_output([efibootmgr], stderr=subprocess.STDOUT, universal_newlines=True)
290
+
# Check the output of `efibootmgr` to find if rEFInd is already installed and present in the boot record
291
+
refind_boot_entry = None
292
+
if matches := re.findall(r'Boot([0-9a-fA-F]{4})\*? rEFInd', efibootmgr_output):
293
+
refind_boot_entry = matches[0]
295
+
# If there's already a Limine entry, replace it
296
+
if refind_boot_entry:
297
+
boot_order = re.findall(r'BootOrder: ((?:[0-9a-fA-F]{4},?)*)', efibootmgr_output)[0]
299
+
efibootmgr_output = subprocess.check_output([
301
+
'-b', refind_boot_entry,
303
+
], stderr=subprocess.STDOUT, universal_newlines=True)
305
+
efibootmgr_output = subprocess.check_output([
308
+
'-b', refind_boot_entry,
310
+
'-p', efi_partition.removeprefix(efi_disk).removeprefix('p'),
311
+
'-l', f'\\efi\\refind\\{boot_file}',
314
+
], stderr=subprocess.STDOUT, universal_newlines=True)
316
+
efibootmgr_output = subprocess.check_output([
320
+
'-p', efi_partition.removeprefix(efi_disk).removeprefix('p'),
321
+
'-l', f'\\efi\\refind\\{boot_file}',
323
+
], stderr=subprocess.STDOUT, universal_newlines=True)
325
+
print("removing unused boot files...")
327
+
if not paths[path]:
331
+
def main() -> None:
333
+
install_bootloader()
335
+
# Since fat32 provides little recovery facilities after a crash,
336
+
# it can leave the system in an unbootable state, when a crash/outage
337
+
# happens shortly after an update. To decrease the likelihood of this
338
+
# event sync the efi filesystem after each update.
339
+
rc = libc.syncfs(os.open(f"{config('efiMountPoint')}", os.O_RDONLY))
341
+
print(f"could not sync {config('efiMountPoint')}: {os.strerror(rc)}", file=sys.stderr)
343
+
if __name__ == '__main__':
pyrox.dev