;;; GNU Guix --- Functional package management for GNU
;;; Copyright © 2024 Lilah Tascheter <lilah@lunabee.space>
;;;
;;; This file is part of GNU Guix.
;;;
;;; GNU Guix is free software; you can redistribute it and/or modify it
;;; under the terms of the GNU General Public License as published by
;;; the Free Software Foundation; either version 3 of the License, or (at
;;; your option) any later version.
;;;
;;; GNU Guix is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;; GNU General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with GNU Guix.  If not, see <http://www.gnu.org/licenses/>.

(define-module (gnu bootloader uki)
  #:use-module (gnu bootloader)
  #:use-module (gnu packages bootloaders)
  #:use-module (gnu packages efi)
  #:use-module (gnu packages linux)
  #:use-module (guix gexp)
  #:use-module (guix modules)
  #:use-module (srfi srfi-1)
  #:export (uefi-uki-bootloader uefi-uki-signed-bootloader))

;; config generator makes script creating uki images
;; install runs script
;; install device is path to uefi dir
(define vendor "Guix")
(define script-loc "/boot/install-uki.scm")

(define* (uefi-uki-configuration-file #:optional cert privkey)
  (lambda* (config entries #:key (old-entries '()) #:allow-other-keys)

    (define (menu-entry->args e)
      (let* ((boot (bootloader-configuration-bootloader config))
             (stub (bootloader-package boot)))
        #~(list "--os-release" #$(menu-entry-label e)
            "--linux" #$(menu-entry-linux e) "--initrd" #$(menu-entry-initrd e)
            "--cmdline" (string-join (list #$@(menu-entry-linux-arguments e)))
            "--stub" #$(file-append stub "/libexec/" (systemd-stub-name))
            #$@(if cert #~("--secureboot-certificate" #$cert) #~())
            #$@(if privkey #~("--secureboot-private-key" #$privkey) #~()))))

    (define (enum-filenames . args) ; same args as iota
      (map (lambda (n) (string-append (number->string n) ".efi"))
        (apply iota (map length args))))

    (program-file "install-uki"
      (with-imported-modules (source-module-closure '((guix build syscalls)
                                                      (guix build utils)))
        #~(let* ((target (cadr (command-line)))
                 (vendir (string-append target "/EFI/" #$vendor))
                 (schema (string-append vendir "/boot.mgr"))
                 (findmnt #$(file-append util-linux "/bin/findmnt"))
                 (efibootmgr #$(file-append efibootmgr "/sbin/efibootmgr")))
            (use-modules (guix build syscalls) (guix build utils)
                         (ice-9 popen) (ice-9 textual-ports))

            (define (out name) (string-append vendir "/" name))
            (define disk
              (call-with-port
                (open-pipe* OPEN_READ findmnt "-fnro" "SOURCE" "-T" target)
                (lambda (port) (get-line port)))) ; only 1 line: the device

            ;; delete all boot entries and files we control
            (when (file-exists? schema)
              (call-with-input-file schema
                (lambda (port)
                  (for-each (lambda (l)
                              (unless (string-null? l)
                                (system* efibootmgr "-B" "-L" l "-q")))
                    (string-split (get-string-all port) #\lf)))))
            (when (directory-exists? vendir) (delete-file-recursively vendir))
            (mkdir-p vendir)

            (define (install port boot? oos)
              (lambda (args label name)
                (let ((minbytes (* 2 (stat:size (stat #$script-loc)))))
                  (put-string port label)
                  (put-char port #\lf)
                  (force-output port) ; make sure space is alloc'd
                  (apply invoke #$(file-append ukify "/bin/ukify")
                    "build" "-o" (out name) args)
                  ;; make sure we have enough space for next install-uki.scm
                  (when (and oos (< (free-disk-space vendir) minbytes)) (oos))
                  (invoke efibootmgr (if boot? "-c" "-C") "-L" label "-d" disk
                    "-l" (string-append "\\EFI\\" #$vendor "\\" name) "-q"))))

            (call-with-output-file schema
              (lambda (port) ; prioritize latest UKIs in limited ESP space
                (for-each (install port #t #f)
                  (list #$@(map-in-order menu-entry->args entries))
                  (list #$@(map-in-order menu-entry-label entries))
                  (list #$@(enum-filenames entries)))
                (for-each ; old-entries can fail (out of space) we don't care
                  (lambda (args label name)
                    (define (cleanup . _) ; do exit early if out of space tho
                      (when (file-exists? (out name)) (delete-file (out name)))
                      (exit))
                    (with-exception-handler cleanup
                      (lambda _ ((install port #f cleanup) args label name))))
                  (list #$@(map-in-order menu-entry->args old-entries))
                  (list #$@(map-in-order menu-entry-label old-entries))
                  (list #$@(enum-filenames old-entries entries))))))))))

(define install-uefi-uki
  #~(lambda (bootloader target mount-point)
      (invoke (string-append mount-point #$script-loc)
              (string-append mount-point target))))

(define* (make-uefi-uki-bootloader #:optional cert privkey)
  (bootloader
    (name 'uefi-uki)
    (package systemd-stub)
    (installer install-uefi-uki)
    (disk-image-installer #f)
    (configuration-file script-loc)
    (configuration-file-generator (uefi-uki-configuration-file cert privkey))))

;; IMPORTANT NOTE: if bootloader install fails, do not turn off your computer! until
;; install succeeds, your system is unbootable.
(define uefi-uki-bootloader (make-uefi-uki-bootloader))
;; use ukify genkey to generate cert and privkey. DO NOT include in store.
(define (uefi-uki-signed-bootloader cert privkey)
  (make-uefi-uki-bootloader cert privkey))