#!/usr/bin/python
# -*- coding: utf-8 -*-

# LTSP-Cluster Account Manager
# Copyright (2008-2009) Stéphane Graber <stgraber@ubuntu.com>, Revolution Linux Inc.
#
# Author: Stéphane Graber <stgraber@ubuntu.com>
#
# 2008, Stéphane Graber <stgraber@ubuntu.com>
# 2009, Stéphane Graber <stgraber@ubuntu.com>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 2 of·
# the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, you can find it on the World Wide
# Web at http://www.gnu.org/copyleft/gpl.html, or write to the Free
# Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston,
# MA 02110-1301, USA.

import pwd, dbus, os, subprocess, syslog, sys, re, socket, thread, base64, time, signal
from Crypto.Cipher import Blowfish

# Get the initial configuration
if len(sys.argv) < 2:
    sys.exit("You must specify a configuration file")

port=None
autologin_root=None
autologin_uidmin=None
autologin_uidmax=None
autologin_groups=None
key=None

try:
    config=open(sys.argv[1],"r")
except:
    sys.exit("Unable to read the configuration file")
groups=[]
for line in config.readlines():
    line=line.strip()
    if line.startswith("port="):
        port=int(line.replace("port=",""))
    elif line.startswith("autologin_root="):
        autologin_root=line.replace("autologin_root=","")
    elif line.startswith("autologin_uidmin="):
        autologin_uidmin=int(line.replace("autologin_uidmin=",""))
    elif line.startswith("autologin_uidmax="):
        autologin_uidmax=int(line.replace("autologin_uidmax=",""))
    elif line.startswith("autologin_groups="):
        autologin_groups=line.replace("autologin_groups=","")
    elif line.startswith("key="):
        key=line.replace("key=","")
if not port or not autologin_root or not autologin_uidmin or not autologin_uidmax or not autologin_groups or not key:
    sys.exit("Configuration file is incomplete") 
config.close()

# Initialize blowfish
cipher = Blowfish.new(key)


# Initialize the syslog logging interface
syslog.openlog("ltsp-cluster-accountmanager")

def Log(msg):
    syslog.syslog(msg)
    print msg

def getUserStatus(uid):
    """Returns the uid, username, home directory, status 
       and Xauthority existance of the user"""
    try:
        user=pwd.getpwuid(uid)

        if os.path.isdir('/run/systemd/users'):
            # ask logind
            status = os.path.exists('/run/systemd/users/%i' % uid)
        else:
            # ask ConsoleKit
            bus = dbus.SystemBus()
            try:
                ck_obj = bus.get_object('org.freedesktop.ConsoleKit', '/org/freedesktop/ConsoleKit/Manager')
            except:
                # If consolekit fails, it'll be spawned and then will work the second time.
                ck_obj = bus.get_object('org.freedesktop.ConsoleKit', '/org/freedesktop/ConsoleKit/Manager')
            ck = dbus.Interface(ck_obj, 'org.freedesktop.ConsoleKit.Manager')
            if ck.GetSessionsForUnixUser(uid):
                status=True
            else:
                status=False

        xauthority=os.path.exists(user.pw_dir+"/.Xauthority")
        kill=os.path.exists(user.pw_dir+"/.ltsp-cluster-accountmanager.kill")
        return {
                 'uid': uid,
                 'name': user.pw_name,
                 'dir': user.pw_dir,
                 'status': status,
                 'xauthority': xauthority,
                 'kill': kill
               }
    except KeyError:
        return {
                 'uid': None,
                 'name': None,
                 'dir': None,
                 'status': None,
                 'xauthority': None,
                 'kill': None
               }

def getProcess():
    "Returns a list of dictionaries containing information on running process"
    proc=[]
    for file in os.listdir("/proc"):
        if file.isdigit():
            try:
                process=open("/proc/"+file+"/status","r")
            except:
                #If we are unable to open the status file it's because the process died in between, just go to the next one
                continue
            for line in process.readlines():
                if line.startswith("Uid"):
                    for uid in line.strip().split("\t")[1:]:
                         proc.append({'pid':int(file),'uid':int(uid)})
            process.close()
    return proc

def getProcessForUID(uid):
    "Returns a list of PID for a given UID"
    processes=getProcess()
    pidlist=set()
    for process in processes:
        if process['uid'] == uid:
            pidlist.add(int(process['pid']))
    return list(pidlist)

def getUIDList():
    "Returns a list of unique uid for all current processus"
    processes=getProcess()
    uidlist=set()
    for process in processes:
        uidlist.add(int(process['uid']))
    return list(uidlist)

def getAutologinUserList():
    "Returns a list of all autologin uid"
    puid=set()
    passwd=open("/etc/passwd","r")
    for line in passwd.readlines():
        if autologin_root in line:
            puid.add(int(line.split(":")[2]))
    passwd.close()
    return list(puid)

def CreateAutologinUser(user, data):
    Log("Creating autologin user \""+user['name']+"\"")
    useradd=subprocess.Popen(["useradd","-K","UID_MIN="+str(autologin_uidmin),"-K","UID_MAX="+str(autologin_uidmax),"-d",autologin_root+user['name'],"-m","-s","/bin/bash","-c","LTSP","-G",autologin_groups,user['name']],stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    if useradd.wait() != 0:
        Log("Failed to create user \""+user['name']+"\"")
        return
    os.mkdir(autologin_root+user['name']+"/.ssh")
    sshkey=open(autologin_root+user['name']+"/.ssh/authorized_keys","w+")
    sshkey.write(data)
    sshkey.close()
    userchown=subprocess.Popen(["chown","-R",user['name']+"."+user['name'],autologin_root+user['name']],stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    userchown.wait()
    runScripts("A",user)

def UmountUserFS(user):
    mounts=open("/proc/mounts","r")
    for line in mounts.readlines():
        if user['dir'] in line:
            umount=subprocess.Popen(["umount",line.split(" ")[1]],stdout=subprocess.PIPE,stderr=subprocess.PIPE)
            umount.wait()
    mounts.close()

def CleanRemainingSession(user):
    "Kill the session of disconnected user"
    Log("Killing user \""+user['name']+"\"")
    UmountUserFS(user)

    # First send a sig term to all process
    for pid in getProcessForUID(user['uid']):
        os.kill(pid,signal.SIGTERM)

    # Then send a sig kill to the remaining ones
    for pid in getProcessForUID(user['uid']):
        try:
            os.kill(pid,signal.SIGKILL)
        except:
            #process already died, just go to the next one
            continue

    runScripts("K",user)

def CleanAutologinSession(user):
    "Remove all unused autologin accounts"
    Log("Removing autologin user \""+user['name']+"\"")
    userdel=subprocess.Popen(["userdel", user['name'], "-r", "-f"],stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    userdel.wait()
    groupdel=subprocess.Popen(["groupdel", user['name']],stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    groupdel.wait()
    userrm=subprocess.Popen(["rm","-Rf", user['dir']],stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    userrm.wait()
    runScripts("D",user)

def CleanSessions():
    "Make sure the system remains clean"
    while 1:
        for uid in getUIDList():
            # Look only for non-system users
            if uid >= 1000 and uid != 65534:
                user=getUserStatus(uid)
                # Get only disconnected users
                if not user['status']:
                    CleanRemainingSession(user)

        for uid in getAutologinUserList():
            user=getUserStatus(uid)
            if not user['status'] and user['xauthority']:
                CleanAutologinSession(user)
        time.sleep(30)

def KillUser(username,admin = False):
    "Make sure the user has the right to kill the session"
    try:
        user=getUserStatus(pwd.getpwnam(username).pw_uid)
    except:
        return False
    if user['status'] and (user['kill'] or admin):
        CleanRemainingSession(user)
    else:
        return False

def adminHandler(connection,command):
    try:
        command=cipher.decrypt(command).strip().replace("\0","")
        if command.startswith("kill "):
            username=command.replace("kill ","")
            KillUser(username,True)
        elif command.startswith("plist "):
            username=command.replace("plist ","")
            plist=subprocess.Popen(["ps","-U",username,"-u",username,"-o","pid,cmd","kstart_time"],stdout=subprocess.PIPE,stderr=None)
            plist.wait()
            for line in plist.stdout.readlines():
                connection.send(line)
    except:
        pass

def clientHandler(connection,addr):
    data = connection.recv(1024).strip()
    account_request = re.findall("ssh-rsa (.*) root@ltsp",data)
    address=addr[0]
    if account_request:
        user={'name': base64.b64encode(address).replace("=","")}
        try:
            user=getUserStatus(pwd.getpwnam(user['name']).pw_uid)
            CleanRemainingSession(user)
            CleanAutologinSession(user)
        except:
            pass
        CreateAutologinUser(user,data)
        connection.send("hello "+user['name']+" "+address+"\n")
    elif data == "username":
        connection.send("hello "+base64.b64encode(address).replace("=","")+" "+address+"\n")
    elif data == "getkey":
        connection.send(key)
    elif data.startswith("kill "):
        username=data.replace("kill ","")
        KillUser(username)
    elif data.startswith("admin "):
        command=data.replace("admin ","").strip()
        adminHandler(connection,command)
    connection.close()
    return

def runScripts(action,user):
    "Execute all scripts for a determined action"
    for file in os.listdir("/usr/share/ltsp-cluster-accountmanager/scripts/"):
        if file.startswith(action):
            environ=os.environ
            environ['LTSP_UID'] = user['uid']
            environ['LTSP_NAME'] = user['name']
            environ['LTSP_DIR'] = user['dir']
            Log("Running: "+file+" for "+user['name'])
            subprocess.Popen([file],stdout=None,stderr=None,env=environ)

# Load the SSH public key
if not os.path.exists("/etc/ssh/ssh_host_rsa_key.pub"):
    sys.exit("Unable to open /etc/ssh/ssh_host_rsa_key.pub")
keyfile=open("/etc/ssh/ssh_host_rsa_key.pub","r")
keyparam=keyfile.read().split(" ")
key=keyparam[0]+" "+keyparam[1]
keyfile.close()

# Prepare the server socket
server=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
server.bind(("0.0.0.0",port))
server.listen(2)

Log("Starting with port="+str(port)+", autologin_uidmin="+str(autologin_uidmin)+", autologin_uidmax="+str(autologin_uidmax)+", autologin_root="+autologin_root)

# Start the cleanup script
thread.start_new_thread(CleanSessions, ())

# Start the multi-thread server
try:
    while 1:
        client,addr=server.accept()
        thread.start_new_thread(clientHandler, (client, addr))
except KeyboardInterrupt:
    server.shutdown(socket.SHUT_RDWR)
    server.close()
    sys.exit(0)
