sommelierd

#!/usr/bin/env ruby
require 'fileutils'
require 'securerandom'

$LOGFILE = (ENV['SOMMELIER_LOG'].to_s.empty?) ? '/tmp/sommelier.log' : ENV['SOMMELIER_LOG']

# set XDG_RUNTIME_DIR if not set
ENV['XDG_RUNTIME_DIR'] = (ENV['XDG_RUNTIME_DIR'].to_s.empty?) ? '/var/run/chrome' : ENV['XDG_RUNTIME_DIR']

SOMMELIER_X_DISPLAY = (ENV['SOMMELIER_X_DISPLAY'].to_s.empty?) ? ':0' : ENV['SOMMELIER_X_DISPLAY']
SOMMELIER_WL_DISPLAY = (ENV['SOMMELIER_WL_DISPLAY'].to_s.empty?) ? 'wayland-1' : ENV['SOMMELIER_WL_DISPLAY']

HELP = <<EOT
sommelierd: Sommelier daemon manager

Usage: sommelierd <command>

  --start     start sommelier daemon
  --stop      stop sommelier daemon
  --restart   restart sommelier daemon
  --log       print daemon output, check and print new outputs repeatedly (blocking)
  --help, -h  print this message

EOT

trap('INT') { abort 'Interrupted' }

def start_somm_x_daemon (display = SOMMELIER_X_DISPLAY)
  Process.fork do
    Process.setproctitle 'sommelier X11 daemon (wrapper script)'

    Process.daemon(false, true)

    # redirect output to log file
    STDOUT.reopen($LOGIO)
    STDERR.reopen($LOGIO)

    # write output to log file directly without saving in buffer
    STDOUT.sync = true
    STDERR.sync = true

    begin
      socket_file = "/tmp/.X11-unix/X#{display.tr(':', '')}"
      IO.popen([ 'sommelier', '-X', '--glamor',
        "--x-display=#{display}",
        "--x-auth=/home/chronos/.Xauthority",
        '--no-exit-with-child',
        'bash', '-c', <<~CMD
          set -eu -o pipefail

          export DISPLAY="#{display}"

          echo -n >> /home/chronos/.Xauthority

          xauth -f /home/chronos/.Xauthority add ${DISPLAY} . "#{SecureRandom.hex(16)}"
          source #{__dir__}/../etc/sommelierrc
        CMD
      ], err: [:child, :out]) do |io|
        puts "[somm_x11]: process running with PID #{io.pid}"
        File.write('/tmp/sommelier-x.pid', io.pid) # write pid to /tmp/sommelier-x.pid

        # write current display number to /tmp/sommelier-x.display (used for checking available display in sommelier.env)
        File.write('/tmp/sommelier-x.display', display)

        print "[somm_x11]: #{io.gets}" until io.eof?
        io.close
      end
      puts '[somm_x11]: process exited'
    ensure
      FileUtils.rm_f [ socket_file, *Dir['/tmp/sommelier-x.{display,pid}'] ]
      puts "[somm_x11]: socket #{socket_file} removed."
    end
  end
end

def start_somm_wl_daemon (display = SOMMELIER_WL_DISPLAY)
  Process.fork do
    Process.setproctitle 'sommelier Wayland daemon (wrapper script)'

    Process.daemon(false, true)

    # redirect output to log file
    STDOUT.reopen($LOGIO)
    STDERR.reopen($LOGIO)

    # write output to log file directly without staying in buffer
    STDOUT.sync = true
    STDERR.sync = true

    begin
      socket_file = "#{ENV['XDG_RUNTIME_DIR']}/#{display}"
      IO.popen([ 'sommelier', '--parent', "--socket=#{display}" ], err: [:child, :out]) do |io|
        puts "[somm_wl]: process running with PID #{io.pid}"
        File.write('/tmp/sommelier-wl.pid', io.pid) # write pid to /tmp/sommelier-wl.pid

        # write current display number to /tmp/sommelier-wl.display (used for checking available display in sommelier.env)
        File.write('/tmp/sommelier-wl.display', display)

        print "[somm_wl]: #{io.gets}" until io.eof?
        io.close
      end
      puts '[somm_wl]: process exited'
    ensure
      FileUtils.rm_f [ socket_file, "#{socket_file}.lock", *Dir['/tmp/sommelier-wl.{display,pid}'] ]
      puts "[somm_wl]: socket #{socket_file} removed."
    end
  end
end

def stop_somm
  # stop_somm: kill sommelier process listed in /tmp/sommelier-*.pid files
  x_pid = ( File.exist?('/tmp/sommelier-x.pid') ) ? File.read('/tmp/sommelier-x.pid') : nil
  wl_pid = ( File.exist?('/tmp/sommelier-wl.pid') ) ? File.read('/tmp/sommelier-wl.pid') : nil

  [ x_pid, wl_pid ].reject(&:nil?).each do |pid|
    begin
      Process.kill('TERM', pid.to_i)
    rescue Errno::ESRCH
    end
  end

  sleep 0.5 # ensure the daemon processes have been killed
  puts "sommelier stopped"
end

def vaild_daemon_running?
  # vaild_daemon_running: check sommelier process validity
  if File.exist?('/tmp/sommelier-x.pid') and File.exist?('/tmp/sommelier-wl.pid')
    x_pid = File.read('/tmp/sommelier-x.pid')
    wl_pid = File.read('/tmp/sommelier-wl.pid')

    [ x_pid, wl_pid ].each do |pid|
      begin
        # use Process.kill to test pid listed in /tmp/sommelier-*.pid files
        Process.kill(0, pid.to_i)
      rescue Errno::ESRCH
        return false
      end
    end
    return true
  else
    return false
  end
end

if $0 =~ /(re)?startsommelier$/ or ARGV[0] =~ /^--(re)?start$/
  stop_somm if $0 =~ /restartsommelier$/ or ARGV[0] == '--restart'

  unless vaild_daemon_running?
    $LOGIO = File.open($LOGFILE, 'w')
    $LOGIO.sync = true
    start_somm_x_daemon
    start_somm_wl_daemon
  end
elsif $0 =~ /stopsommelier/ or ARGV[0] == '--stop'
  stop_somm
elsif ARGV[0] == '--log'
  io = File.open($LOGFILE, 'r'); sleep(0.2) until File.exist?($LOGFILE)
  # print daemon output, check and print new outputs repeatedly
  while (sleep 0.2; log = io.read) do
    puts log.gsub(/^(\[.*?\])/, "\e[0;33m\\1\e[0m") unless log.empty?
  end
elsif ARGV[0] == '--status'
  abort "sommelier not running, see #{$LOGFILE} for more info." unless vaild_daemon_running?
  puts "sommelier process running: #{File.read('/tmp/sommelier-x.pid')} #{File.read('/tmp/sommelier-wl.pid')}"
elsif [ '--help', '-h' ].include?(ARGV[0])
  puts HELP
else
  STDERR.puts <<~EOT
    Invaild command: #{ARGV[0]}
    
    Type `sommelierd --help` for usage.
  EOT
end