;;;
;;; Lightweight VPS image builder, by Peter Bex.
;;;
;;; This program is hereby placed in the public domain.
;;;
;;
;; This provides a way of building custom images to run on VPSes which
;; can automatically configure themselves, without requiring
;; heavyweight, difficult to install solutions like Hashicorp's
;; "Packer" (which has no Debian package) and/or CloudInit (which in
;; Debian requires not one, but two versions of Python to be
;; installed!)
;;
;; Heavily inspired by Debian's "build-openstack-debian-image" shell
;; script, provided by the "openstack-debian-images" package.
;;
;; Currently only Debian Jessie and CloudStack are supported.
;; Auto-resizing is NOT supported at the moment, you'll need to know
;; in advance how big the VM's root disk will be.  In fact, you'll
;; need to hack this script to change what it does.
;;
;; Maybe later, if I use this script more, it can be turned into
;; a proper customisable program/library.
;;
;; To run, this requires the following packges:
;; qemu-utils, kpartx, parted, mbr
;;
;; It depends on the "scsh-process" egg
;;
;; Before running it, you can create a "users" directory containing
;; the public keys of all the users you want to prepopulate the image
;; with.  The names of these files need to be USER:GROUP1,...  so for
;; example, to create a user "peter" who is a member of groups "sudo"
;; and "adm", copy his id_rsa.pub file to "users/peter:sudo,adm".

(module vps-builder ()

  (import chicken scheme)
  (use data-structures extras files posix scsh-process)

  (define debug? #t)

  (reset-handler (lambda () #f))
  
  ;; Might it be useful to put this in scsh-process?
  (define-syntax run*
    (syntax-rules ()
      ((_ ?pf ?redir ...)
       (begin
         (when debug?
           (fprintf (current-error-port) "$ ~S\n" `?pf))
         (receive (status normal? pid)
             (run ?pf ?redir ...)
           (unless (and normal? (zero? status))
             (error (sprintf "Pipeline ~S exited abnormally, with code: ~A"
                      `(run ?pf ?redir ...) status))))))))

  (define minimal-packages
    `("sudo" "locales" "extlinux" "openssh-server" "file" "kbd"
      ;; TODO: What about other architectures?
      "linux-image-amd64"))

  ;; It's necessary to be root due to chroot and mount calls.
  ;; TODO: Maybe support using sudo instead?  For now, just call
  ;; the whole script with sudo.  Otherwise we'll need to mess
  ;; with PATH and so on, too.
  (define (check-privileges)
    (unless (zero? (current-effective-user-id))
      (error "Sorry, you must be root to build an image!")))

  (define (check-users)
    (when (or (not (directory? "users")) (null? (glob "users/*:*")))
      (error (conc "You need to put public ssh keys in the "
                   "\"users\" directory, named like "
                   "USERNAME:GROUP1,..."))))

  (define (create-image-file raw-file size-in-gb)
    (let ((size (sprintf "~AG" size-in-gb)))
      (run* (qemu-img create ,raw-file ,size))))

  (define (convert-to-qcow2 raw-file qcow-file)
    ;; TODO: Figure out if compat option is really needed
    (run* (qemu-img convert -c -f raw -O qcow2 -o compat=0.10
                    ,raw-file ,qcow-file)))

  (define (partition-image image-file)
    (run* (parted -s ,image-file mktable msdos))
    (run* (parted -s ,image-file -a optimal
             mkpart primary ext3 1Mi 100%))
    (run* (parted -s ,image-file set 1 boot on))
    (run* (install-mbr ,image-file)))

  (define (prepare-image raw-image size-in-gb)
    (create-image-file raw-image size-in-gb)
    (partition-image raw-image))

  (define (call-with-devmapped-image image-file proc)
    (let ((kpartx-output #f))
      (dynamic-wind
          (lambda ()
            ;; TODO: run/string doesn't (and can't) check exit status!
            (set! kpartx-output (run/string (kpartx -asv ,image-file))))
          (lambda ()
            ;; Yeah this is stupid...
            (let* ((parts (string-split kpartx-output " "))
                   (loopback-dev (list-ref parts 2)))
              (proc (make-pathname "/dev/mapper" loopback-dev))))
          (lambda ()
            (set! kpartx-output #f)
            ;; For some reason kpartx will sometimes(?) fail
            (handle-exceptions exn #f
              (run* (kpartx -d ,image-file)))))))

  (define (prepare-filesystem device)
    ;; Apparently, operating on ext2 is much faster than ext3, so we
    ;; make it ext2, and convert to ext3 later
    (run* (mkfs.ext2 ,device)))

  (define (finalize-filesystem device)
    ;; Add journal, turning it into ext3
    (run* (tune2fs -j ,device)))

  (define (call-with-mounted-device device proc)
    (let ((tempdir #f))
      (dynamic-wind
          (lambda ()
            (set! tempdir (create-temporary-directory))
            (run* (mount -o loop ,device ,tempdir)))
          (lambda () (proc tempdir))
          (lambda ()
            (run* (umount ,tempdir))
            (delete-directory tempdir #t)
            (set! tempdir #f)))))

  ;; Create a very basic setup which allows running programs inside
  ;; a chroot (mount essential filesystems).
  (define (with-running-system root-dir thunk)
    (dynamic-wind
        (lambda ()
          (run* (chroot ,root-dir mount /proc))
          ;; Don't run daemons in the chroot upon apt-get install
          (install-file root-dir "assets/chroot-helpers/policy-rc.d"
                        "/usr/sbin/policy-rc.d" "root" "root" #o755)
          (setenv "DEBIAN_FRONTEND" "noninteractive"))
        (lambda () (thunk))
        (lambda ()
          (unsetenv "DEBIAN_FRONTEND")
          (delete-file* (make-pathname root-dir "/usr/sbin/policy-rc.d"))
          (run* (chroot ,root-dir umount /proc)))))

  (define (install-basic-system target-dir package-list)
    (let ((include (sprintf "--include=~A"
                     (string-intersperse package-list ","))))
      (run* (debootstrap
             --verbose ,include "jessie" ,target-dir
             "http://ftp.surfnet.nl/os/Linux/distr/debian"))))

  (define (install-file root-dir source target owner group mode)
    (let* ((full-path (make-pathname root-dir target))
           (full-path (if (directory? full-path)
                       (make-pathname full-path
                                      (pathname-strip-directory source))
                       full-path)))
      (file-copy source full-path #t) ; allow clobber
      ;; Must run in chroot to use correct passwd/group db
      (run* (chroot ,root-dir chown ,(conc owner ":" group) ,target))
      (change-file-mode full-path mode)))

  (define (install-directory root-dir target owner group mode)
    (let ((full-path (make-pathname root-dir target)))
      (create-directory full-path)
      ;; Must run in chroot to use correct passwd/group db
      (run* (chroot ,root-dir chown ,(conc owner ":" group) ,target))
      (change-file-mode full-path mode)))

  (define (install-packages root-dir . packages)
    (run* (chroot ,root-dir
                  apt-get install -y
                  -o Dpkg::Options=--force-confdef
                  -o Dpkg::Options=--force-confold
                  ,@packages)))

  (define (configure-basic-system root-dir)
    ;;;; Configure apt, FS and disable console bleeping (just in case)
    (install-file root-dir "assets/package-manager/apt/apt-settings"
                  "/etc/apt/apt.conf.d/90custom-config"
                  "root" "root" #o644)
    (install-file root-dir "assets/fstab"
                  "/etc/fstab" "root" "root" #o644)
    (install-file root-dir "assets/modprobe/blacklist.conf"
                  "/etc/modprobe.d/blacklist.conf" "root" "root" #o644)

    ;;;; Setup network and related settings
    (delete-file* (make-pathname
                   root-dir "etc/udev/rules.d/70-persistent-net.rules"))
    (install-file root-dir "assets/network/interfaces"
                  "/etc/network/interfaces" "root" "root" #o644)
    (install-file root-dir "assets/network/cloudstack-guest-setup"
                  "/etc/dhcp/dhclient-exit-hooks.d" "root" "root" #o644))

  (define (update-packages root-dir)
    (run* (chroot ,root-dir apt-get update))
    (run* (chroot ,root-dir apt-get upgrade -y
                  -o Dpkg::Options=--force-confdef
                  -o Dpkg::Options=--force-confold))
    (run* (chroot ,root-dir apt-get autoremove -y))
    (run* (chroot ,root-dir apt-get clean))
    (run* (chroot ,root-dir apt-get autoclean)))

  (define (make-bootable root-dir)

    (define (rooted-glob dir pattern)
      (let ((f (car (glob (make-pathname `(,root-dir ,dir) pattern)))))
        (make-pathname `("/" ,dir) (pathname-strip-directory f))))

    (let* ((kernel (rooted-glob "boot" "vmlinuz-*"))
           (ramdisk (rooted-glob "boot" "initrd.img-*"))
           (template (with-input-from-file
                         "assets/boot/extlinux.conf" read-string))
           (conf (string-translate* template
                                    `(("{KERNEL}" . ,kernel)
                                      ("{RAMDISK}" . ,ramdisk)))))
      (with-output-to-file (make-pathname root-dir "extlinux.conf")
        (lambda () (write-string conf)))
      (run* (extlinux --install ,root-dir))))
  
  (define (setup-firewall root-dir)
    (install-packages root-dir "ferm" "sshguard")
    (install-file root-dir "assets/firewall/ferm.conf"
                  "/etc/ferm/ferm.conf" "root" "adm" #o644))

  ;; No proper monitoring solution yet, but at least vnstat is useful
  ;; for keeping an eye on traffic even if we don't automate it yet.
  (define (setup-monitoring root-dir)
    (install-packages root-dir "vnstat"))

  ;; Create a more convenient default environment.  This is highly
  ;; dependent on taste so you might want to change this.
  (define (customize-environment root-dir)
    (install-file root-dir "assets/zsh/zshrc" "/etc/skel/.zshrc"
                  "root" "root" #o644)
    (install-file root-dir "assets/default/useradd"
                  "/etc/default/useradd" "root" "root" #o644)
    (install-packages root-dir "zsh"))


  ;; Create user and copy matching users/*:* file to .ssh/authorized_keys
  (define (create-users root-dir)
    (for-each
     (lambda (pubkey)
       (let* ((fn (pathname-strip-directory pubkey))
              (user+cs-groups (string-split fn  ":"))
              (user (car user+cs-groups))
              (cs-groups (cadr user+cs-groups))
              (~ (make-pathname `("/" "home") user))
              (~/.ssh (make-pathname ~ ".ssh")))
         ;; Use -p to set empty *crypted* password.  This ensures the
         ;; user has no password (which differs from having an empty one!)
         ;; and can choose to set a password.  If we didn't supply -p, the
         ;; password would be *locked*, which means "passwd" will prompt
         ;; for a password, but there's none, so it can't be changed.
         (run* (chroot ,root-dir useradd -p "" -m -G ,cs-groups ,user))
         (install-directory root-dir ~/.ssh user user #o700)
         (install-file root-dir pubkey
                       (make-pathname ~/.ssh "authorized_keys")
                       user user #o600)))
     (glob "users/*:*")))

  (define (build-image image-base-name size-in-gb)
    (let ((raw-image (make-pathname '() image-base-name ".raw"))
          (qcow-image (make-pathname '() image-base-name ".qcow2")))
      (check-privileges)
      (check-users)
      (prepare-image raw-image size-in-gb)

      (call-with-devmapped-image
       raw-image
       (lambda (dev)
         (prepare-filesystem dev)

         (call-with-mounted-device
          dev
          (lambda (mountpoint)
            (install-basic-system mountpoint minimal-packages)
            (configure-basic-system mountpoint)

            (with-running-system
             mountpoint
             (lambda ()
               (update-packages mountpoint)
               (make-bootable mountpoint)

               (setup-firewall mountpoint)

               (setup-monitoring mountpoint)

               (customize-environment mountpoint)
               (create-users mountpoint)))))

         (finalize-filesystem dev)))

      (convert-to-qcow2 raw-image qcow-image)
      (printf "Success!  Images finished, raw: ~A, qcow2: ~A\n"
        raw-image qcow-image)))

;; Just do it!
(handle-exceptions exn (signal exn) ; re-throw, to unwind dynamic extents
  (build-image "debian-jessie" 20))
)