Add utmp support. Fix #39

butterfly/routes.py

@@ -125,13 +125,14 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
     terminals = set()
 
     def pty(self):
+        self.determine_user()
         self.pid, self.fd = pty.fork()
         if self.pid == 0:
             self.shell()
         else:
             self.communicate()
 
-    def shell(self):
+    def determine_user(self):
         if self.callee is None and (
                 tornado.options.options.unsecure and
                 tornado.options.options.login):
@@ -149,6 +150,7 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
 
         assert self.callee is not None
 
+    def shell(self):
         try:
             os.chdir(self.path or self.callee.dir)
         except:
@@ -232,6 +234,10 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
     def communicate(self):
         fcntl.fcntl(self.fd, fcntl.F_SETFL, os.O_NONBLOCK)
 
+        utils.add_user_info(
+            self.fd, self.pid,
+            self.callee.name, self.request.headers['Host'])
+
         def utf8_error(e):
             self.log.error(e)
 
@@ -330,10 +336,10 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
         if events & ioloop.ERROR:
             self.log.info('Error on fd %d, closing' % fd)
             # Terminated
-            self.on_close()
             self.close()
 
     def on_close(self):
+        utils.rm_user_info(self.fd, self.pid, self.callee.name)
         if self.fd is not None:
             self.log.info('Closing fd %d' % self.fd)
 

butterfly/utils.py

@@ -18,7 +18,11 @@
 
 import os
 import pwd
+import time
+import sys
+import struct
 from logging import getLogger
+from collections import namedtuple
 import subprocess
 import re
 
@@ -222,3 +226,117 @@ def get_socket_env(inode):
                                         key, val = keyval.split('=', 1)
                                         env[key] = val
                                 return env
+
+
+utmp_struct = struct.Struct('hi32s4s32s256shhiii4i20s')
+
+
+if sys.version_info[0] == 2:
+    b = lambda x: x
+else:
+    def b(x):
+        if isinstance(x, str):
+            return x.encode('utf-8')
+        return x
+
+
+def get_utmp_file():
+    for file in (
+            '/var/run/utmp',
+            '/var/adm/utmp',
+            '/var/adm/utmpx',
+            '/etc/utmp',
+            '/etc/utmpx',
+            '/var/run/utx.active'):
+        if os.path.exists(file):
+            return file
+
+
+def get_wtmp_file():
+    for file in (
+            '/var/log/wtmp',
+            '/var/adm/wtmp',
+            '/var/adm/wtmpx',
+            '/var/run/utx.log'):
+        if os.path.exists(file):
+            return file
+
+UTmp = namedtuple(
+            'UTmp',
+            ['type', 'pid', 'line', 'id', 'user', 'host',
+             'exit0', 'exit1', 'session',
+             'sec', 'usec', 'addr0', 'addr1', 'addr2', 'addr3', 'unused'])
+
+
+def utmp_line(type, pid, fd, user, host, ts):
+    return UTmp(
+        type,  # Type, 7 : user process
+        pid,  # pid
+        b('pts/%d' % fd),  # line
+        b('/%d' % fd),  # id
+        b(user),  # user
+        b(host),  # host
+        0,  # exit 0
+        0,  # exit 1
+        0,  # session
+        int(ts),  # sec
+        int(10 ** 6 * (ts - int(ts))),  # usec
+        0,  # addr 0
+        0,  # addr 1
+        0,  # addr 2
+        0,  # addr 3
+        b('')  # unused
+    )
+
+
+def add_user_info(fd, pid, user, host):
+    utmp = utmp_line(7, pid, fd, user, host, time.time())
+    for kind, file in {
+            'utmp': get_utmp_file(),
+            'wtmp': get_wtmp_file()}.items():
+        if not file:
+            continue
+        try:
+            with open(file, 'rb+') as f:
+                s = f.read(utmp_struct.size)
+                while s:
+                    entry = UTmp(*utmp_struct.unpack(s))
+                    if kind == 'utmp' and entry.id.rstrip(b'\0') == utmp.id:
+                        # Same id recycling
+                        f.seek(f.tell() - utmp_struct.size)
+                        f.write(utmp_struct.pack(*utmp))
+                        break
+                    s = f.read(utmp_struct.size)
+                else:
+                    f.write(utmp_struct.pack(*utmp))
+        except Exception:
+            log.warning('Unable to write utmp info to ' + file, exc_info=True)
+
+
+def rm_user_info(fd, pid, user):
+    import traceback
+    log.warning(traceback.format_stack())
+    utmp = utmp_line(8, pid, fd, user, '', time.time())
+    for kind, file in {
+            'utmp': get_utmp_file(),
+            'wtmp': get_wtmp_file()}.items():
+        if not file:
+            continue
+        try:
+            with open(file, 'rb+') as f:
+                s = f.read(utmp_struct.size)
+                while s:
+                    entry = UTmp(*utmp_struct.unpack(s))
+                    if kind == 'utmp' and entry.id.rstrip(b'\0') == utmp.id:
+                        log.warning('Found Writing ' + str(utmp) + ' ' + kind)
+                        # Same id closing
+                        f.seek(f.tell() - utmp_struct.size)
+                        f.write(utmp_struct.pack(*utmp))
+                        break
+                    s = f.read(utmp_struct.size)
+                else:
+                    log.warning('Else Writing ' + str(utmp) + ' ' + kind)
+                    f.write(utmp_struct.pack(*utmp))
+
+        except Exception:
+            log.warning('Unable to update utmp info to ' + file, exc_info=True)