#!/usr/bin/env ruby

require "socket"
require "fileutils"
require "json"
require "shellwords"

SOCKET_PATH = "/run/malpd/malpd.sock"

def human_capacity(bytes)
  units = ["B", "KB", "MB", "GB", "TB", "PB"]
  i = 0
  size = bytes.to_f
  while size >= 1000 && i < units.length - 1
    size /= 1000
    i += 1
  end
  precision = i <= 3 ? 0 : 1
  "%g %s" % [size.round(precision), units[i]]
end

def drive_type(name)
  return "nvme" if name.start_with?("nvme")
  rotational = File.read("/sys/block/#{name}/queue/rotational").strip
  rotational == "1" ? "hdd" : "ssd"
end

def read_smart(name)
  out = `smartctl -j -A -H /dev/#{name} 2>/dev/null`
  return nil if out.empty?
  data = JSON.parse(out) rescue nil
  return nil unless data

  result = {}
  passed = data.dig("smart_status", "passed")
  result["passed"] = passed unless passed.nil?

  if (t = data.dig("temperature", "current"))
    result["temperature"] = t
  end
  if (poh = data.dig("power_on_time", "hours"))
    result["power_on_hours"] = poh
  end

  if name.start_with?("nvme")
    log = data["nvme_smart_health_information_log"] || {}
    result["percentage_used"] = log["percentage_used"] if log.key?("percentage_used")
    result["media_errors"]    = log["media_errors"]    if log.key?("media_errors")
    result["critical_warning"] = log["critical_warning"] if log.key?("critical_warning")
  else
    attrs = data.dig("ata_smart_attributes", "table") || []
    attrs.each do |a|
      raw = a.dig("raw", "value")
      next if raw.nil?
      case a["id"]
      when 5   then result["reallocated_sectors"] = raw
      when 10  then result["spin_retry"] = raw
      when 197 then result["pending_sectors"] = raw
      when 199 then result["crc_errors"] = raw
      end
    end
  end

  result
end

def read_inodes(mount)
  out = `stat -f -c '%c %d' #{Shellwords.escape(mount)} 2>/dev/null`.strip
  return nil if out.empty?
  total, free = out.split.map(&:to_i)
  return nil unless total && total > 0
  [total, total - free]
end

def read_filesystems
  out = `findmnt -J -b -o SOURCE,TARGET,FSTYPE,SIZE,USED,OPTIONS 2>/dev/null`
  return [] if out.empty?
  data = JSON.parse(out) rescue nil
  return [] unless data

  skip_fstypes = %w[tmpfs devtmpfs efivarfs]
  result = []
  seen = {}
  stack = Array(data["filesystems"]).dup
  until stack.empty?
    node = stack.shift
    stack.concat(node["children"]) if node["children"]
    source = node["source"].to_s
    fstype = node["fstype"].to_s
    target = node["target"].to_s
    next unless source.start_with?("/dev/")
    next if skip_fstypes.include?(fstype)
    size = node["size"].to_i
    next if size <= 0
    key = [source, target]
    next if seen[key]
    seen[key] = true

    opts = node["options"].to_s.split(",")
    entry = {
      "type"     => "filesystem",
      "source"   => source,
      "mount"    => target,
      "fstype"   => fstype,
      "size"     => size,
      "used"     => node["used"].to_i,
      "readonly" => opts.include?("ro")
    }
    if (ino = read_inodes(target))
      entry["inode_total"] = ino[0]
      entry["inode_used"]  = ino[1]
    end
    result << entry
  end
  result
end

def parse_mem(s)
  return 0 unless s
  n, unit = s.split
  mult = case unit
         when "KiB" then 1024
         when "MiB" then 1024 ** 2
         when "GiB" then 1024 ** 3
         when "TiB" then 1024 ** 4
         else 1
         end
  n.to_i * mult
end

def read_vm_dominfo(name)
  out = `virsh --connect qemu:///system dominfo #{Shellwords.escape(name)} 2>/dev/null`
  info = {}
  out.each_line do |line|
    key, value = line.split(":", 2).map { |s| s&.strip }
    next unless key && value
    case key
    when "State"       then info["state"] = value
    when "CPU(s)"      then info["vcpus"] = value.to_i
    when "Max memory"  then info["memory"] = parse_mem(value)
    when "Used memory" then info["used_memory"] = parse_mem(value)
    when "Autostart"   then info["autostart"] = (value == "enable")
    end
  end
  info
end

def read_vm_info(name)
  ssh_out = `ssh -Ti /root/.ssh/malp-vm-key malp@#{name}`
  JSON.parse(ssh_out)
end

def read_virtual_machines
  list_out = `virsh --connect qemu:///system list --all --name 2>/dev/null`
  return [] unless $?.success?
  names = list_out.split("\n").map(&:strip).reject(&:empty?)
  return [] if names.empty?

  names.map do |name|
    dom_info = read_vm_dominfo(name)
    vm_info = read_vm_info(name)
    {
      "type"        => "vm",
      "name"        => name,
    }.merge(dom_info).merge(vm_info)
  end
end

def read_ups
  out = `apcaccess 2>/dev/null`
  return nil unless $?.success? && !out.empty?

  raw = {}
  out.each_line do |line|
    next unless line.include?(":")
    key, value = line.split(":", 2).map(&:strip)
    raw[key] = value
  end

  result = { "type" => "ups" }
  result["model"] = raw["MODEL"] if raw["MODEL"]
  result["status"] = raw["STATUS"] if raw["STATUS"]

  if raw["BCHARGE"]
    result["charge"] = raw["BCHARGE"].to_f
  end
  if raw["TIMELEFT"]
    result["timeleft"] = raw["TIMELEFT"].to_f
  end
  if raw["LINEV"]
    result["line_voltage"] = raw["LINEV"].to_f
  end
  if raw["LOADPCT"]
    result["load"] = raw["LOADPCT"].to_f
  end
  if raw["NOMBATTV"]
    result["battery_voltage_nom"] = raw["NOMBATTV"].to_f
  end
  if raw["BATTV"]
    result["battery_voltage"] = raw["BATTV"].to_f
  end

  result
end

def handle_info(conn)
  info = []

  Dir.glob("/sys/block/*").each do |block_path|
    name = File.basename(block_path)
    next unless name.match?(/^(sd[a-z]+|nvme\d+n\d+)$/)

    device_path = File.realpath("#{block_path}/device") rescue next
    next if device_path.include?("/usb")

    model = (File.read("#{block_path}/device/model").strip rescue "Unknown")
    bytes = File.read("#{block_path}/size").strip.to_i * 512

    entry = {
      "type" => "drive",
      "drivetype" => drive_type(name),
      "name" => name,
      "model" => model,
      "capacity" => human_capacity(bytes)
    }
    if (smart = read_smart(name))
      entry["smart"] = smart
    end
    info << entry
  end

  info.concat(read_filesystems)
  info.concat(read_virtual_machines)

  if (ups = read_ups)
    info << ups
  end

  if conn
    conn.puts JSON.generate(info)
  else
    puts JSON.pretty_generate(info)
  end
end

if ARGV.include?("-t")
  handle_info(nil)
  exit
end

if Process.uid != 0
  $stderr.puts "malpd must be run as root"
  exit(1)
end

if ENV["LISTEN_FDS"]
  server = UNIXServer.for_fd(3)
else
  FileUtils.mkdir_p(File.dirname(SOCKET_PATH))
  FileUtils.rm_f(SOCKET_PATH)

  server = UNIXServer.new(SOCKET_PATH)
  File.chmod(0660, SOCKET_PATH)
  FileUtils.chown(nil, "apache", SOCKET_PATH)
end

Signal.trap("TERM") { server.close }
Signal.trap("INT") { server.close }

loop do
  conn = server.accept
  request = conn.gets&.chomp

  begin
    case request
    when nil
      # empty request
    when "info"
      handle_info(conn)
    else
      conn.puts "unknown request: #{request}"
    end
  rescue => e
    $stderr.puts "error handling request: #{e.message}"
  ensure
    conn.close
  end
rescue IOError
  break
end