malpd: report drive info

bin/malpd

@@ -2,9 +2,61 @@
 
 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)
@@ -20,8 +72,6 @@ FileUtils.chown(nil, "apache", SOCKET_PATH)
 Signal.trap("TERM") { server.close }
 Signal.trap("INT") { server.close }
 
-$stdout.puts "malpd listening on #{SOCKET_PATH}"
-
 loop do
   conn = server.accept
   request = conn.gets&.chomp
@@ -30,6 +80,8 @@ loop do
     case request
     when nil
       # empty request
+    when "info"
+      handle_info(conn)
     else
       conn.puts "unknown request: #{request}"
     end