(define-module (vkraus services simple-firewall)
  #:use-module ((gnu services) #:select (service-type service-extension))
  #:use-module ((gnu packages linux) #:select (nftables))
  #:use-module (gnu services shepherd)
  #:use-module (guix gexp)
  #:use-module (guix packages)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-9 gnu)
  #:use-module (ice-9 match)
  #:declarative? #t
  #:export
  (
   <simple-firewall-configuration>
   simple-firewall-configuration
   simple-firewall-configuration?
   nftables-package opened-tcp-ports opened-udp-ports
   set-nftables-package set-opened-tcp-ports set-opened-udp-ports
   merge-firewall-configurations
   simple-firewall-service-type
   ))

(define-immutable-record-type <simple-firewall-configuration>
  (simple-firewall-configuration nft tcp-ports udp-ports)
  simple-firewall-configuration?
  (nft nftables-package set-nftables-package)
  (tcp-ports opened-tcp-ports set-opened-tcp-ports)
  (udp-ports opened-udp-ports set-opened-udp-ports))

(define (append-if-not-present item list)
  (if (member item list)
      list
      `(,@list ,item)))

(define (remove-duplicates all)
  (let append ((lst all)
               (unique '()))
    (match lst
      (() unique)
      ((x lst ...)
       (append lst (append-if-not-present x unique))))))

(define merge-firewall-configurations
  (match-lambda*
    (() (simple-firewall-configuration #f '() '()))
    ((($ <simple-firewall-configuration> nft tcp udp)
      cfg ...)
     (match (apply merge-firewall-configurations cfg)
       (($ <simple-firewall-configuration> nfts tcps udps)
        (simple-firewall-configuration
         (if (and nft nfts)
             (error "The nftables package is overriden twice. Please choose.")
             (or nft nfts))
         (remove-duplicates `(,@tcp ,@tcps))
         (remove-duplicates `(,@udp ,@udps))))))))

(define port->string
  (match-lambda
    ((? symbol? high-level-name)
     (with-output-to-string
       (lambda () (display high-level-name))))
    ((? integer? number)
     (with-output-to-string
       (lambda () (display number))))))

(define simple-firewall-shepherd-service
  (match-lambda
    (($ <simple-firewall-configuration> nft tcp udp)
     (let ((config-file
            (plain-file
             "simple-firewall-configuration"
             (format #f
                     "flush ruleset

table inet firewall {
    chain inbound {
    	# By default, drop all traffic unless it meets a filter
    	# criteria specified by the rules that follow below.
        type filter hook input priority 0; policy drop;

        # Allow traffic from established and related packets.
        ct state established,related accept

        # Drop invalid packets.
        ct state invalid drop

        # Allow loopback traffic.
        iifname lo accept

        # Allow all ICMP and IGMP traffic, but enforce a rate limit
        # to help prevent some types of flood attacks.
        ip protocol icmp limit rate 4/second accept
        meta l4proto ipv6-icmp limit rate 4/second accept
        ip protocol igmp limit rate 4/second accept

        # Allow DHCPv6 incoming replies
        ip6 daddr fe80::/64 udp sport 547 udp dport 546 ct state { new, untracked } accept

        # Allow TCP ports
        ~atcp dport { ~a } accept

        # Allow UDP ports
        ~audp dport { ~a } accept
    }

    chain forward {
        # Drop everything (assumes this device is not a router)
        type filter hook forward priority 0; policy drop;

    }

    chain outbound {
        # Allow all outbound traffic
        type filter hook output priority 0; policy accept;
    }
}
"
		     (if (null? tcp)
			 "# "
			 "")
                     (string-join
                      (map port->string tcp)
                      ", ")
		     (if (null? udp)
			 "# "
			 "")
                     (string-join
                      (map port->string udp)
                      ", "))))
           (nft-package
            (or nft nftables)))
       (shepherd-service
        (documentation
         "Simple firewall that only allows inbound packets to specific ports")
        (provision '(simple-firewall))
        (start
         #~(lambda _
             (invoke #$(file-append nft-package "/sbin/nft")
                     "--file" #$config-file)))
        (stop
         #~(lambda _
             (invoke #$(file-append nft-package "/sbin/nft")
                     "flush" "ruleset"))))))))

(define simple-firewall-service-type
  (service-type
   (name 'simple-firewall)
   (extensions
    (list
     (service-extension
      shepherd-root-service-type
      (lambda (cfg)
        `(,(simple-firewall-shepherd-service cfg))))))
   (compose
    (lambda (extensions)
      (apply append extensions)))
   (extend
    (lambda (a bs)
      (apply merge-firewall-configurations a bs)))
   (description
    "A very simple firewall service with some ports opened.")))