#!/usr/bin/env ruby

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

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 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

    info << {
      "type" => "drive",
      "drivetype" => drive_type(name),
      "model" => model,
      "capacity" => human_capacity(bytes)
    }
  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

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)

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