(define-module (tw services wireguard)
  #:use-module (ice-9 format)
  #:use-module (ice-9 match)
  #:use-module (ice-9 regex)
  #:use-module ((srfi srfi-1) #:select (append-map every))
  #:use-module ((srfi srfi-26) #:select (cut))
  #:use-module (gnu services)
  #:use-module (gnu services base)
  #:use-module (gnu services configuration)
  #:use-module (gnu services vpn)
  #:use-module (guix gexp)
  #:use-module ((guix records) #:select (match-record))
  #:use-module ((guix utils) #:select (current-source-directory))
  #:use-module (tw services secrets)
  #:export (%wireguard-peers
            tw-wireguard-configuration
            tw-wireguard-service-type))

(define %wireguard-peers
  ;; Order in the following list is significant! It determines what IPs are assigned.
  (let ((peers '(("lap" "lap/DvCb8xXLUCqcaPEx8kCRcoeV4ScTMVZW5hvvNzA=" #f)
                 ("lud" "lud/9sbXVdOYXxOkRgAB+b/17QxbwllfJY/pbA3/MkE=" 58921)
                 ("vin" "vin/Im+sOszZFE01UF1+QlyxLP1PsPXJgTz4KmgvL3Y=" 58921)
                 ("fp4" "fp4/aLAVBADTy+UGmNh011w1CFOOwq70Df6EWlZRkAs=" #f)
                 ("pi3" "pi3/ThUH4qDTuyvNQIiiyy2dbziF/xLRTwO0+vcUoVY=" 58922)
                 ("frm" "frm/YGu1BfXUl4jrN0PTFMNdTQXWPSuY1wEpz5W9C2Y=" #f)
                 ("btl" "btl/kAgD+DVXsApNn53JCZdgZ9iJvVpFZVpa3Z+rrj4=" 58923))))

    (map (match-lambda*
           ((i (short-name public-key port))
            (cons (string-append short-name ".twilken.net")
                  (wireguard-peer
                   (name (string-append short-name ".wg"))
                   (endpoint (and port (format #f "~a.twilken.net:~d" short-name port)))
                   (public-key public-key)
                   (preshared-key (string-append "/etc/wireguard/" short-name ".psk"))
                   (allowed-ips (list (format #f "10.0.0.~d/32" (+ i 1))
                                      (format #f "fc00::~d/128" (+ i 1)))))))
           (args (error "Unknown peer spec" args)))
         (iota (length peers)) peers)))

(define (wireguard-peers-list? object)
  (and (list? object)
       (every (compose string? car) object)
       (every (compose wireguard-peer? cdr) object)))

(define-configuration/no-serialization tw-wireguard-configuration
  (this-host
   (string)
   "The host name of the machine being configured.")
  (peers
   (wireguard-peers-list %wireguard-peers)
   "An alist of WireGuard peers to install.")
  (private-key-file
   (string "/etc/wireguard/private.key")
   "Where to store this host's private key."))

(define (other-peers this-host peers)
  (delq (assoc-ref peers this-host) (map cdr peers)))

(define (tw-wireguard-service config)
  "Create a full WireGuard config from the personal network CONFIG."
  (match-record config <tw-wireguard-configuration> (this-host peers private-key-file)
    (unless (assoc this-host peers)
      (error "No peer config found for host" this-host))
    (match-record (assoc-ref peers this-host) (@@ (gnu services vpn) <wireguard-peer>) (endpoint allowed-ips)
      (wireguard-configuration
       (addresses
        (map (lambda (cidr)
               (let ((ipv4 (string-match "/32$" cidr))
                     (ipv6 (string-match "/128$" cidr)))
                 (cond
                  (ipv4 (regexp-substitute #f ipv4 'pre "/24"))
                  (ipv6 (regexp-substitute #f ipv6 'pre "/64"))
                  (#t cidr))))
             allowed-ips))
       (port
        (if endpoint
            (string->number (cadr (string-split endpoint #\:)))
            58921))
       (private-key private-key-file)
       (peers (other-peers this-host peers))))))

(define (cut-string-at-char str char-pred)
  "Return the first part of STR up to the first occurrence of CHAR-PRED."
  (substring str 0 (string-index str char-pred)))

(define (peer->ips peer)
  "Extract IP addresses assigned to the given `wireguard-peer' PEER."
  (map (cut cut-string-at-char <> #\/)
       (wireguard-peer-allowed-ips peer)))

(define (tw-wireguard-hosts config)
  "Generate a hosts file entries from the personal WireGuard network CONFIG."
  (define (peer->entries peer)
    (map (cut host <> (wireguard-peer-name peer))
         (peer->ips peer)))
  (append-map (compose peer->entries cdr)
              (tw-wireguard-configuration-peers config)))

(define (tw-wireguard-secrets config)
  "Install secrets for the host's private key and preshared keys with peers."
  (define (local-file-here path)
    (local-file
     (canonicalize-path
      (string-append
       (current-source-directory) "/" path))))
  (match-record config <tw-wireguard-configuration> (this-host peers private-key-file)
    (define short-host (cut-string-at-char this-host #\.))
    (define private-key
      (secret
       (encrypted-file
        (local-file-here (string-append "files/wireguard/" short-host ".key.enc")))
       (destination private-key-file)))
    (define (peer->secret peer)
      (let ((short-peer (cut-string-at-char (wireguard-peer-name peer) #\.)))
        (secret
         (encrypted-file
          (local-file-here
           (string-append "files/wireguard/" short-host "-" short-peer ".psk.enc")))
         (destination
          (string-append "/etc/wireguard/" short-peer ".psk")))))
    (cons private-key (map peer->secret (other-peers this-host peers)))))

(define tw-wireguard-service-type
  (service-type
   (name 'tw-wireguard)
   (description "Set up my personal WireGuard network.")
   (extensions
    (cons* (service-extension hosts-service-type tw-wireguard-hosts)
           (service-extension secrets-service-type tw-wireguard-secrets)
           ;; FIXME: `wireguard-service-type' cannot be extended, so copy its
           ;; service-extensions directly.
           (map (lambda (ext)
                  (service-extension (service-extension-target ext)
                                     (compose (service-extension-compute ext)
                                              tw-wireguard-service)))
                (service-type-extensions wireguard-service-type))))))