#!/usr/bin/env python2.4

import sys
import os
import subprocess
import errors

# TODO: minimise shell/process calls

def execute(*args, **kwargs):
  if isinstance(args, tuple):
    args = list(args)
  PIPE = subprocess.PIPE
  p = subprocess.Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE, **kwargs)
  p.wait()
  errors = p.stderr.read()
  if errors:
    raise OSError("Error executing %s:\n%s" % (" ".join(args), errors))
  return p.stdout.read()

def exec_dcop(*args):
  args = [str(arg) for arg in args]
  print "calling dcop with %r" % args
  return execute("dcop", *args, **{"env": {"DISPLAY": ":0.0", "HOME": "/home/davidf"}}).split("\n")

try:
  # using native python dcop may be slightly faster
  import pcop
  def dcop(*args):
    if len(args) == 0:
      return pcop.app_list()
    elif len(args) == 1:
      return pcop.obj_list(*args)
    elif len(args) >= 3:
      return str(pcop.dcop_call(args[0], args[1], args[2], args[3:]))
    else:
      return exec_dcop(*args)
except ImportError:
  dcop = exec_dcop

def shell(*args):
  return execute(*args, **{"shell": True})

def list_konsole_processes():
  processlist = shell("ps --no-heading -o uid,pid -C konsole")
  for processline in processlist.split("\n"):
    if not processline: continue
    uid, pid = processline.strip().split()
    uid = uid.strip()
    if not uid.isdigit():
      continue
    pid = pid.strip()
    if not pid.isdigit():
      continue
    yield int(uid), int(pid)

class Process:
  """OS process analyzer"""
  def __init__(self, pid):
    if isinstance(pid, (str, unicode)):
      self.pid = int(pid)
    else:
      self.pid = pid
    self.cmdline, self.cwd, self.exe, self.env, self.fds, self.children = None, None, None, {}, [], {}
    self.get_proc_info()
    self.uid = shell("ps --no-heading -o '%%u' %d" % self.pid).strip()
    self.get_children()

  def safe_read(self, filename):
    """returns contents of file, but won't raise an error on IO problems (e.g. permissions)"""
    try:
      return open(filename).read()
    except Exception, e:
      return ""

  def safe_readlink(self, filename):
    """returns contents of file, but won't raise an error on IO problems (e.g. permissions)"""
    try:
      return os.readlink(filename)
    except Exception, e:
      return ""

  def safe_listdir(self, dirname):
    """returns contents of file, but won't raise an error on IO problems (e.g. permissions)"""
    try:
      return os.listdir(dirname)
    except Exception, e:
      return []

  def get_proc_info(self):
    """reads information about the process from /proc"""
    procdir = os.path.join("/proc", str(self.pid))
    self.cmdline = self.safe_read(os.path.join(procdir, "cmdline")).replace("\0", " ")
    self.cwd = self.safe_readlink(os.path.join(procdir, "cwd"))
    self.exe = self.safe_readlink(os.path.join(procdir, "exe"))
    environ_items = self.safe_read(os.path.join(procdir, "environ")).split("\0")
    self.env = {}
    for environ_item in environ_items:
      if not environ_item: continue
      if "=" not in environ_item: continue
      variable, value = environ_item.split("=", 1)
      self.env[variable] = value
    fddir = os.path.join(procdir, "fd")
    self.fds = []
    for fd in self.safe_listdir(fddir):
      fullfd = os.path.join(fddir, fd)
      if os.path.islink(fullfd):
        fdfile = self.safe_readlink(fullfd)
        if fdfile is not None:
          self.fds.append(fdfile)

  def get_vim_info(self):
    """gets vim-specific info"""
    swapfiles = {}
    for fd in self.fds:
      basefd = os.path.basename(fd)
      fddir = os.path.dirname(fd)
      if basefd.startswith(".") and basefd[:-1].endswith(".sw"):
        filename = os.path.join(fddir, basefd[1:-4])
        swapfiles[fd] = filename
    return swapfiles

  def get_children(self):
    """finds child process IDs and creates child Process objects in self.children"""
    child_pids = shell("pgrep -P %d" % self.pid).split()
    self.children = {}
    for child_pid in child_pids:
      if not child_pid: continue
      child_pid = int(child_pid)
      self.children[child_pid] = Process(child_pid)

  def plaintree(self, indent=0):
    """prints out a process tree with the given indent"""
    # plaintree = shell("pstree -alupG %s" % session_pid, shell=True)
    print " "*indent + self.cmdline + "[%s@%s]" % (self.uid, self.cwd)
    for child, child_process in self.children.iteritems():
      child_process.plaintree(indent+2)

  def htmltree(self):
    """prints out a process tree with the given indent"""
    html = "<li>" 
    html += "<span title='process %d' class='cmdline'>%s</span>\n" % (self.pid, self.cmdline) 
    html += "    <span class='uid'>%s</span> in <span class='cwd'>%s</span>\n" % (self.uid, self.cwd)
    if os.path.basename(self.exe) == "vim":
      swapfiles = self.get_vim_info()
      if swapfiles:
        html += " editing <ul>\n"
        edititems = []
        for swap, edit in swapfiles.iteritems():
          if swap:
            edititem = "<li class='editfile' title='%s'>%s</li>\n" % (swap, edit)
            edititems.append(edititem)
        html += "".join(edititems) + "</ul>\n"
    html += "</li>\n"
    if self.children:
      html += "<ul>\n"
      for child, child_process in self.children.iteritems():
        html += child_process.htmltree()
      html += "</ul>\n"
    return html

class KonsoleSession:
  """Information on konsole session"""
  def __init__(self, dcopid, sessionid):
    self.dcopid = dcopid
    self.sessionid = sessionid
    self.pid = self.dcop("sessionPID").strip()
    self.name = self.dcop("sessionName").strip()
    self.process = Process(self.pid)

  def dcop(self, *args):
    """passes a dcop command to this konsole session"""
    return dcop(self.dcopid, self.sessionid, *args)

class KonsoleProcess:
  """Information on konsole sessions"""
  def __init__(self, pid):
    if isinstance(pid, (str, unicode)):
      self.pid = int(pid)
    else:
      self.pid = pid
    self.dcopid = "konsole-%d" % self.pid
    self.sessioncount = self.dcop("konsole", "sessionCount").strip()
    self.currentsessionid = self.dcop("konsole", "currentSession").strip()
    self.readsessions()

  def readsessions(self):
    self.sessions = []
    for sessionnum in range(1, int(self.sessioncount)+1):
      sessionid = self.dcop("konsole", "sessionId", sessionnum).strip()
      session = KonsoleSession(self.dcopid, sessionid)
      session.iscurrent = (sessionid == self.currentsessionid)
      self.sessions.append(session)

  def dcop(self, *args):
    """passes a dcop command to this konsole instance"""
    return dcop(self.dcopid, *args)

  def gethtml(self):
    """returns html for the sessions in this konsole"""
    konsoleheader = "Konsole %s: %s sessions" % (self.pid, self.sessioncount)
    html = "<h1>" + konsoleheader + "</h1>"
    for session in self.sessions:
      html += "<h2>" + session.name 
      if getattr(session, "currentsession", False):
        html += " (current)"
      html += "</h2>"
      html += "<ul>"
      html += session.process.htmltree()
      html += "</ul>"
    return html

def xlog(m):
  f = open("/tmp/x", "a")
  f.write(m+"\n")
  f.close()

def iterkonsoles():
  # konsole_pids = [d.replace("konsole-", "", 1) for d in dcop() if d.startswith("konsole-")]
  current_uid = os.getuid()
  konsole_pids = list_konsole_processes()
  xlog("pid: %r" % (konsole_pids,))
  for uid, konsole_pid in konsole_pids:
    xlog("pid %r" % konsole_pid)
    if uid != current_uid:
      print "user id for process %r differs: %r, %r" % (konsole_pid, uid, current_uid)
      os.setuid(uid)
      print "switched uid", os.getuid()
    try:
      dcop("konsole-%s" % konsole_pid)
    except OSError:
      # this is a dead konsole
      continue
    xlog("got %r" % konsole_pid)
    try:
      konsole = KonsoleProcess(konsole_pid)
      xlog("got konsole %r" % konsole_pid)
      yield konsole
    except Exception, e:
      xlog("konsoleprocess make error: %s" % errors.ConsoleErrorHandler().traceback_str())

def plain():
  for konsole in iterkonsoles():
    konsoleheader = "Konsole %s: %s sessions" % (konsole.pid, konsole.sessioncount)
    print konsoleheader
    print "-"*len(konsoleheader)
    for session in konsole.sessions:
      print session.name + ":"
      session.process.plaintree(2)
    print

def main():
  import time
  if "--output" in sys.argv:
    homedir = os.getenv("HOME", None) 
    if not homedir or homedir == "/":
      homedir = "/root"
    basedir = os.path.join(homedir, "konsolelog")
    if not os.path.isdir(basedir):
      os.mkdir(basedir)
    starttime = time.time()
    indexhtml = "<html><head><link rel='stylesheet' href='konsole.css' type='text/css'>"
    indexhtml += "<title>Konsole sessions at %s</title></head>" % time.asctime()
    indexhtml += "<body>"
    # try:
    #   for konsole in iterkonsoles():
    #     pass
    # except Exception, e:
    #   xlog("iterproblem: %s" % str(e))
    for konsole in iterkonsoles():
      xlog("again")
      try:
        html = konsole.gethtml()
        konsolefilename = "konsole-%d-%s.html" % (konsole.pid, time.strftime("%H%M"))
        konsolefile = open(os.path.join(basedir, konsolefilename), "w")
        konsolefile.write("<html><head><link rel='stylesheet' href='konsole.css' type='text/css'>")
        konsolefile.write("<title>Konsole session %d</title></head>" % konsole.pid)
        konsolefile.write("<body><p>Konsole session at %s</p>" % time.asctime())
        konsolefile.write(html)
        konsolefile.write("</body></html>")
        konsolefile.close()
      except Exception, e:
        open("/var/log/cron", "a").write("Exception on konsole: %s\n" % str(e))
        continue
      indexhtml += "<p><a href='%s'>konsole %d</a></p>" % (konsolefilename, konsole.pid)
    endtime = time.time()
    indexhtml += "<p>Generated in %0.2f seconds</p>" % (endtime-starttime)
    indexhtml += "</body></html>"
    indexfile = open(os.path.join(basedir, "index.html"), "w")
    indexfile.write(indexhtml)
    indexfile.close()
  else:
    starttime = time.time()
    print "<html><head>"
    print "<title>Konsole sessions at %s</title>" % time.asctime()
    print "<link rel='stylesheet' href='konsole.css' type='text/css'>"
    print "</head><body>"
    for konsole in iterkonsoles():
      print konsole.gethtml()
    endtime = time.time()
    print "<p>Generated in %0.2f seconds</p>" % (endtime-starttime)
    print "</body></html>"

if __name__ == "__main__":
  if "--html" in sys.argv:
    try:
      # pcop.dcop_call("konsole-6972", "konsole", "sessionCount", ())
      main()
    except Exception, e:
      from jToolkit import errors
      xlog("error: %s" % errors.ConsoleErrorHandler().traceback_str())
  else:
    plain()