#
# ctop.rb 
# written by KAMEZAWA Hiroyuki <kamezawa.hiroyu@jp.fujitsu.com>
# Copyright 2009 Fujitsu Limited
#
# Changelog:
#
# v002
#   - fixed leak of file descriptor
#   - mount/umount <-> reload data problem is fixed.
#   - "mount twice" problem is fixed.
#   - removed R key for reload all. it's now automatic
#   - handle "noprefix" mount option
#   - show mount option in help window
#   - add cpuset support
#   - add command-mode
# v001
#   - first version released
#   - cpu, cpuacct, memory subsys is supported
#   known bugs -> noprefix, umount, mount twice
#
require 'cgroup.rb'
require 'curses'
require 'etc'
require 'timeout'
require 'singleton'

DIRWIN_LINES=7
DIRWIN_FIELDS= DIRWIN_LINES - 2
UPKEY=256
DOWNKEY=257
RIGHTKEY=258
LEFTKEY=259

#mode
SHOWMOUNT=0
SHOWTASKS=1
SHOWSUBSYS=2

#for 'ps'
PID=0
STATE=1
PPID=2
UID=3
COMMAND=4
PGID=5

#for process status filter
RUNNING=0

class Cursor
  def initialize(name)
    @subsysname=name
    @cursor=0
    @mode=SHOWTASKS
    @info_startline=0
    @info_endline=0
    @show_only_running = 0
    @user_name_filter=nil
    @command_name_filter=nil
  end

  def pos
    @cursor
  end

  def mode
    @mode
  end

  def change_mode
    case @mode
    when SHOWTASKS then @mode=SHOWSUBSYS
    when SHOWSUBSYS then @mode=SHOWTASKS
    end
  end
  #
  # Filter for PS-MODE
  #
  def process_status_filter(stat)
    return true if (@show_only_running == 0)
    return true if (stat =="R")
    return false
  end

  def user_name_filter(str)
    return true if (@user_name_filter == nil)
    return true if (@user_name_filter == str)
    return false
  end

  def command_name_filter(str)
    return true if (@command_name_filter == nil)
    return true if (str =~ /#{@command_name_filter}/)
    return false
  end

  def toggle_show_only_running
    if (@show_only_running == 0) then
      @show_only_running = 1
    else
      @show_only_running = 0
    end
  end

  def set_user_name_filter(str)
    @user_name_filter=str
  end

  def set_command_name_filter(str)
    @command_name_filter=str
  end

  #
  # Scroll management for infowin 
  #
  def info_startline
    @info_startline
  end

  def set_infoendline(num)
    @info_endline=num
  end

  def set_infoline(num)
    if ((num < 0) || (num >= @info_endline)) then
      @info_startline=0
    else
      @info_startline=num
    end
  end
  #
  # chdir() for subsys.
  #
  def move(direction)
    subsys =$allsubsys[@subsysname]
    if (subsys == nil) then return
    end
    if (direction == -1) then
      @cursor -= 1 if @cursor > 0
    elsif (direction == 1)
      @cursor += 1 if @cursor < subsys.size-1
    end
  end
end

class Current
  include Singleton
  def initialize
    @index=-1
    @name=nil
    @cursor=nil
    @subsys=nil
    @subsys_cursor = Hash.new
  end

  def set(x)
    @index=x
    if (x == -1) then
      @index, @name, @cursor, @subsys = -1, "help", nil, nil
    else
      @name = $subsys_array[x]
      @subsys = $allsubsys[@name]
      if (@subsys_cursor[@name] == nil) then
        @subsys_cursor[@name] = Cursor.new(@name)
      end
      @cursor = @subsys_cursor[@name]
    end
  end

  def move (dir)
    case dir
    when "left"
      @index -= 1 if (@index > -1)
    when "right"
      @index += 1 if (@index < $subsys_array.size - 1)
    end
    set(@index)
  end

  def name
    @name
  end

  def cursor
    @cursor
  end

  def subsys
    @subsys
  end
end

$cur = Current.instance
#
# A fucntion for change subsys view. $current_subsys is moved.
#
def move_current_subsys(dir)
  $cur.move(dir)
end
#
# Move cursor in directory view
#
def move_current_dir(direction)
  cursor = $cur.cursor
  return  if (cursor == nil)
  cursor.move(direction)  
end


def detect_dirlist_position(subsysname, subsys)
  pos = 0
  size=subsys.size
  cursor = $cur.cursor
  return [0, 0, 0] if cursor == nil

  pos = cursor.pos
  if ((size < 4) || (pos <= 2)) then 
      head=0
      tail=4
  elsif (pos < size - 2) then
      head=pos-1
      tail=pos+2
  else
      head = size - 4
      tail = size - 1
  end
  return [pos, head, tail]
end

def get_owner_name(name)
  stat = File.stat(name)
  begin
    info = Etc::getpwuid(stat.uid)
    uname = info.name
  rescue
    $barwin.addstr($!)
    uname = stat.uid.to_s
  end

  begin
    info = Etc::getgrgid(stat.gid)
    gname = info.name
  rescue
    gname = stat.gid.to_s
  end
  sprintf("\t-\t(%s/%s)", uname, gname)
end

def draw_dirlist(dirwin, subsys, pos, head, tail)
  lines=1
  i=head
  while i <= tail
    name = subsys.ent(i)
    if (name == nil) then break
    end

    dirwin.standout if (i == pos)
    dirwin.setpos(lines, 3)
    dirwin.addstr(name + get_owner_name(name))
    dirwin.standend if (i == pos)
    lines+=1
    i += 1
  end
end

def draw_dirwin(dirwin)
  dirwin.clear
  dirwin.box(?|,?-,?*)
  dirwin.setpos(0, 1)

  -1.upto($subsys_array.size) do |x|
    dirwin.addstr("-")
    if (x == -1) then str="help"
    else str=sprintf("%s",$subsys_array[x])
    end
    break if (str == nil)

    dirwin.standout if (str == $cur.name)
    dirwin.addstr(str)
    dirwin.standend if (str == $cur.name)
  end

  dirwin.setpos(6,dirwin.maxx-32)
  dirwin.addstr("[#{Time.now.asctime}]")
  #
  # Show directory list
  #
  subsys = $cur.subsys
  if subsys != nil then
    #Reload information 
    subsys.reload
    size=subsys.size
    pos = detect_dirlist_position($cur.name, subsys)
    draw_dirlist(dirwin, subsys, pos[0], pos[1], pos[2])
  end
end

#
#
# for infowin
#
#

def draw_infowin_limited(infowin, cursor, data)
  #
  # Generate Header
  #
  str = yield nil # write a header if necessary
  if (str != nil) then
    draw=1
    infowin.setpos(0,2)
    infowin.addstr(str)
  else
    draw=0
  end
  #
  # print a line whici is in the window
  #
  startline = cursor.info_startline
  endline = cursor.info_startline + infowin.maxy-2
  startline.upto(endline) do |linenumber|
    x = data.at(linenumber)
    return if (x == nil) #no more data
    str = yield(x)
    infowin.setpos(draw, 2)
    infowin.addstr(str)
    draw = 1+infowin.cury
    break if (draw == infowin.maxy)
  end

  cursor.set_infoendline(data.size)  
end


def show_mount_info(infowin)
  if ($allsubsys.empty?) then
    $barwin.addstr("cgroups are not mounted\n")
  end
  $allsubsys.each do |name, subsys|
    formated = sprintf("%12s\t%s\t#%s\n",
                       name, subsys.mount_point, subsys.option)
    infowin.addstr(formated)
  end
  #$barwin.addstr("mounted subsystems")
  #
  # Help
  #
  infowin.addstr("Command\n")
  infowin.addstr("[LEFT, RIGHT]\t move subsystems\n")
  infowin.addstr("[UP, DOWN]\t move directory\n")
  infowin.addstr("[n, b]\t\t scorll information window\n")
  infowin.addstr("[s]\t\t switch shown information (ps-mode/stat-mode)\n")
  infowin.addstr("[r]\t\t set refresh rate\n")
  infowin.addstr("[c]\t\t Enter command-mode\n")

  infowin.addstr("ps mode option\n")
  infowin.addstr("[t]\t\t (ps-mode)toggle show only running process\n")
  infowin.addstr("[u]\t\t (ps-mode)set/unset user name filter\n")
  infowin.addstr("[f]\t\t (ps-mode)set/unset command name filter")

end

#
# Read /proc/<pid>/status file and fill data[] array, return it
#
def parse_pid_status(f, es)
  input = f.readline
  input =~ es
  return $1
end

def parse_process(pid)
  #
  # Status
  #
  data = Array.new
  stat = nil

  stat = catch(:bad_task_status) do
    data[PID]=pid.to_i
    begin
      f = File.open("/proc/#{pid}/status", "r")

      #Name
      data[COMMAND] = parse_pid_status(f,/^Name:\s+(.+)/)
      #State
      data[STATE] = parse_pid_status(f, /^State:\s+([A-Z]).+/)
      # TGID: Is thread grouo leader ?
      if (parse_pid_status(f, /^Tgid:\s+(.+)/) != pid) then
        throw :bad_task_status, false
      end
      #skip PID
      input = f.readline
      #PPID
      data[PPID]= parse_pid_status(f,/^PPid:\s+(.+)/)
      ppid=data[PPID]
      #TracerPID
      input = f.readline 
      #UID
      uid = parse_pid_status(f,/^Uid:\s+([0-9]+).+/)
      begin 
        info=Etc::getpwuid(uid.to_i)
        data[UID]=info.name
      rescue 
        data[UID]=uid
      end
    rescue
      throw :bad_task_status, false 
    ensure
      f.close unless f.nil?
    end
  end
  return data unless stat.nil?
  return nil
end

#
# PS-MODE
# Cat "tasks" file and visit all /proc/<pid>/status file
# All information will be pushed into "ps" array
#
def show_tasks(subsys, cursor, infowin)
  # Get Name of Current Cgroup and read task file
  ps = Array.new  
  catch :quit do
    group = subsys.ent(cursor.pos)
    throw :quit,"nogroup" if group==nil
    tasks = subsys.tasks(group)
    throw :quit,"nogroup" if tasks==nil
     
    tasks.each do |x|
      data = parse_process(x)
      next if (data == nil)
      next unless (cursor.process_status_filter(data[STATE]))
      next unless (cursor.command_name_filter(data[COMMAND]))
      ps.push(data) if (cursor.user_name_filter(data[UID]))
    end
    #
    # Sort ps's result, "R" first.
    #
    ps.sort! do |x , y|
      if (x[STATE] == "R" && y[STATE] != "R") then
        -1
      elsif (x[STATE] != "R" && y[STATE] == "R") then
        1
      else
        0
      end
    end
  end
  
  return if (ps.size == 0)

  draw_infowin_limited(infowin, cursor, ps)do |x|
    if (x == nil) then
      sprintf("%6s %6s %8s %5s %16s", "PID","PPID","USER","STATE", "COMMAND")
    else
      sprintf("%6d %6d %8s %5s %16s",
              x[PID], x[PPID], x[UID], x[STATE], x[COMMAND])
    end
  end

  unless ($cur.cursor.process_status_filter("S")) then
    $barwin.addstr("[r]")
  end
  unless ($cur.cursor.user_name_filter("badnamemandab")) then
    $barwin.addstr("[u]")
  end
  unless ($cur.cursor.command_name_filter("badnamemandab")) then
    $barwin.addstr("[c]")
  end
end

def show_subsys_stat(subsys, cursor, infowin)
  group = subsys.ent(cursor.pos)
  return if group == nil
  data = subsys.stat(group)
  return if data == nil
  draw_infowin_limited(infowin, cursor, data) do |x|
    next if x == nil
    if (x[0].size > 24) then
      len = x[0].size - 24
      x[0].slice!(0..len)
    end
    sprintf("%24s\t%s", x[0], x[1])
  end
end


def draw_infowin(infowin)
  infowin.clear
  cursor = $cur.cursor
  if cursor == nil then
    mode = SHOWMOUNT
  else
    mode = cursor.mode
  end
  #
  # If no subsys is specified, just show mount information.
  #

  case mode
  when SHOWMOUNT
    show_mount_info(infowin)
  when SHOWTASKS
    $barwin.addstr("[ps-mode]")
    show_tasks($cur.subsys, cursor, infowin)
  when SHOWSUBSYS
    $barwin.addstr("[stat-mode]")
    show_subsys_stat($cur.subsys, cursor, infowin)
  end
end

def set_scroll(infowin, direction)
  cursor = $cur.cursor
  return if (cursor == nil)

  if (direction == 1) then 
    curline=cursor.info_startline
    cursor.set_infoline(curline+infowin.maxy)
  else
    curline=cursor.info_startline
    cursor.set_infoline(curline-infowin.maxy)
  end
end

def change_mode
    if ($cur.cursor != nil) then    
      $cur.cursor.change_mode
    end
end

def toggle_running_filter(ch)
  if ($cur.cursor != nil) then
    $cur.cursor.toggle_show_only_running
  end
end

def hit_any_key(str, window)
  window.addstr(str) if str != nil
  window.addstr("\n[Hit Any Key]")
  window.getch
end


def user_name_filter(infowin)
  infowin.clear
  infowin.refresh
  x=infowin.maxx-6
  y=3
  window = Curses::stdscr.subwin(y, x, 8, 3)
  window.box(?|,?-,?*)
  window.setpos(1,1)
  window.addstr("user name filter:")
  window.refresh
  str=window.getstr
  window.close
  return if (str == nil)
  cursor= $cur.cursor
  str = nil if (str.size==0)
  cursor.set_user_name_filter(str) if (cursor != nil)
end

def command_name_filter(infowin)
  infowin.clear
  infowin.refresh
  x=infowin.maxx-6
  y=3
  window = Curses::stdscr.subwin(y, x, 8, 3)
  window.box(?|,?-,?*)
  window.setpos(1,1)
  window.addstr("command name filter:")
  window.refresh
  str=window.getstr
  window.close
  return if (str == nil)
  cursor= $cur.cursor
  str = nil if (str.size==0)
  cursor.set_command_name_filter(str) if (cursor != nil)
end

def set_refresh_time(infowin)
  infowin.clear
  infowin.refresh
  x=infowin.maxx-6
  y=3
  window = Curses::stdscr.subwin(y, x, 8, 3)
  window.box(?|,?-,?*)
  window.setpos(1,1)
  window.addstr("set refresh time:")
  window.refresh
  str=window.getstr
  window.close
  return nil if (str == nil)
  return str.to_i
end

def smart_print(str, window)
  if (window.maxx - window.curx < str.size-2) then
    window.addstr("\n"+str)
  else
    window.addstr(str)
  end
end

def show_writable_files(subsys, cursor, infowin)
  group = subsys.ent(cursor.pos)
  return nil if group == nil
  data = subsys.writable_files(group)
  return nil if data == nil
  ent=1
  data.sort!
  data.each do |x|
    str = sprintf("%2d: %s ", ent, File.basename(x))
    ent=ent+1
    smart_print(str, infowin)
  end
  infowin.refresh
  return data
end

def chown_all_files(uid, gid, group, infowin)
  Dir.foreach(group) do |x|
    infowin.addstr(x)
    name = group+"/"+x
    if (x == ".") then
      begin
        File.chown(uid, gid, name)
     rescue
        hit_any_key("Error:"+$!, infowin)
        break
     end
    end
    next if File.directory?(name)
    begin
      File.chown(nil, gid, name)
    rescue
      hit_any_key("Error:"+$!, infowin)
      break
    end
  end
end

def check_mkrmdir_string(str, infowin)
  if (str =~ /\//) then
    infowin.addstr("don't include /\n")
    return false
  end
  return true
end

def command_mode(infowin)
  return if ($cur.subsys == nil)
  infowin.clear
  data = show_writable_files($cur.subsys, $cur.cursor, infowin)
  return if data==nil
  $barwin.clear
  $barwin.addstr("[command-mode]")
  $barwin.refresh
  infowin.addstr("\n")
  smart_print("96: chown(UID) ", infowin)
  smart_print("97: chown(GID) ", infowin)
  smart_print("98: mkdir ", infowin)
  smart_print("99: rmdir",infowin)
  infowin.addstr("\n\nEdit which ?[and Hit return]:")
  endline = infowin.cury+1
  str=infowin.getstr
  name=nil
  group = $cur.subsys.ent($cur.cursor.pos)

  case str.to_i
  when 1..95
    catch :select_failure do
      val = str.to_i
      infowin.setpos(0, 0)
      name = data.at(val - 1)
      throw :select_failure if name==nil
      #get input
      infowin.setpos(endline, 0)
      str = sprintf("#echo to >%s:", File.basename(name))
      infowin.addstr(str)
      infowin.refresh
      str = infowin.getstr
      #write
      begin
        f = File.open(name, "w") {|f| f.write(str) }
      rescue
        hit_any_key("Error:"+$!, infowin)
      end
    end
  #chown (UID)
  when 96
    infowin.addstr("change owner id of all files to:")
    str = infowin.getstr
    if (str =~ /\D/) then
      begin
        info = Etc::getgrnam(str)
        uid = info.gid
      rescue
        hit_any_key("Error:"+$!, infowin)
        uid=nil
      end
    else
      uid = info.to_i
    end
    chown_all_files(uid, -1, group, infowin)
  #chown(GID)
  when 97
    infowin.addstr("change group id of all files to:")
    str = infowin.getstr
    if (str =~ /\D/) then
      begin
        info = Etc::getgrnam(str)
        gid = info.gid
      rescue
        hit_any_key("Error:"+$!, infowin)
        gid=nil
      end
    else
      gid = info.to_i
    end
    chown_all_files(-1, gid, group, infowin)
  #mkdir
  when 98
    infowin.addstr("mkdir -.enter name:")
    str = infowin.getstr
    if (check_mkrmdir_string(str, infowin)) then
       begin
         Dir.mkdir(group+"/"+str)
         hit_any_key("Error:"+$!, infowin)
       rescue
         hit_any_key("Error:"+$!, infowin)
       end
    else
       hit_any_key(nil, infowin)
    end
  #rmdir
  when 99
    infowin.addstr("rmdir -.enter name:")
    str = infowin.getstr
    if (check_mkrmdir_string(str, infowin)) then
       begin
         Dir.rmdir(group+"/"+str)
         hit_any_key("Error:"+$!, infowin)
       rescue
         hit_any_key("Error:"+$!, infowin)
       end
    else
       hit_any_key(nil, infowin)
    end
  end
  $barwin.clear
end


#
# Main loop is following
#

#
# For stdscreen
#
# Check /proc/mounts and read all subsys.
#
refresh_all

#
# Main loop. create windows and wait for inputs
#
Curses::init_screen
begin
  $lines=Curses::lines
  $cols=Curses::cols
  off=0
  #
  # Create window
  #
  dirwin = Curses::stdscr.subwin(DIRWIN_LINES, $cols, off, 0)
  #for misc info
  off+=DIRWIN_LINES
  $barwin = Curses::stdscr.subwin(1, $cols, off, 0);
  $barwin.standout
  off+=1
  infowin = Curses::stdscr.subwin($lines-off, $cols, off, 0)
  mode=SHOWTASKS  
  quit=0
  refresh_time=15
  
  while quit == 0 
    #$barwin.clear

    #$barwin.addstr("Info:")
    draw_dirwin(dirwin)
    draw_infowin(infowin)
    dirwin.refresh
    infowin.refresh
    $barwin.refresh
    #
    # handle input.
    # 
    $barwin.clear
    Curses::setpos(0,0)
    ch=0
    begin
      Timeout::timeout(refresh_time) do
        ch=Curses::getch
      end
    rescue Timeout::Error
      #$barwin.addstr("timeout")
    end

    #check espace sequence
    if ch == 27 then
      ch = Curses::getch
      if ch == 91 then
        ch = Curses::getch
        case ch
          when 65 then ch = UPKEY
          when 66 then ch = DOWNKEY
          when 67 then ch = RIGHTKEY
          when 68 then ch = LEFTKEY
        end
      end
    end
    #
    #
    #
    if (check_and_refresh_mount_info) then
      $cur.set(-1)
    end
   
    #str=sprintf("%d",ch)
    #$barwin.addstr(Time.now.asctime)
    case ch
      when ?q
        quit=1
        break
      when LEFTKEY then move_current_subsys("left")
      when RIGHTKEY then move_current_subsys("right")
      when UPKEY then move_current_dir(-1)
      when DOWNKEY then move_current_dir(1)
      when ?s then change_mode
      when ?n then set_scroll(infowin, 1)
      when ?b then set_scroll(infowin, -1)
      when ?t then toggle_running_filter(nil)
      when ?u then user_name_filter(infowin)
      when ?f then command_name_filter(infowin)
      when ?c then command_mode(infowin)
      when ?r
        newval = refresh_time=set_refresh_time(infowin)
        refresh_time = newval if newval != nil
    end
  end
ensure
  Curses::close_screen
end