/neatx/lib/app/nxnode.py
http://neatx.googlecode.com/ · Python · 307 lines · 230 code · 17 blank · 60 comment · 3 complexity · cee8c2a33a01a67805d5d42d25d5665f MD5 · raw file
- #
- #
- # Copyright (C) 2007 Google Inc.
- #
- # 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, write to the Free Software
- # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
- # 02110-1301, USA.
- """nxnode program.
- This program is started once per session. It receives commands from nxserver
- via a Unix socket and updates the session database based on nxagent's output.
- """
- import logging
- import pwd
- import select
- import signal
- import socket
- import sys
- import gobject
- from neatx import cli
- from neatx import constants
- from neatx import daemon
- from neatx import errors
- from neatx import node
- from neatx import serializer
- from neatx import session
- from neatx import utils
- PROGRAM = "nxnode"
- _SESSION_START_TIMEOUT = 30
- class NxNodeContext(object):
- def __init__(self):
- self.cfg = None
- self.uid = None
- self.username = None
- self.sessrunner = None
- self.sessid = None
- self.session = None
- self.sessmgr = None
- self.processes = None
- def _GetUserUid(username):
- return pwd.getpwnam(username).pw_uid
- def ValidateRequest(req):
- if not (isinstance(req, dict) and
- node.REQ_FIELD_CMD in req and
- node.REQ_FIELD_ARGS in req):
- raise errors.GenericError("Incomplete request")
- class ClientOperations(object):
- def __init__(self, ctx):
- self._ctx = ctx
- def __call__(self, cmd, args):
- logging.info("Received request: %r, %r", cmd, args)
- if cmd == node.CMD_STARTSESSION:
- return self._StartSession(args)
- elif cmd == node.CMD_ATTACHSESSION:
- assert len(args) == 2
- return self._AttachSession(args[0], args[1])
- elif cmd == node.CMD_RESTORESESSION:
- return self._RestoreSession(args)
- elif cmd == node.CMD_TERMINATESESSION:
- return self._TerminateSession(args)
- elif cmd == node.CMD_GET_SHADOW_COOKIE:
- return self._GetShadowCookie()
- else:
- raise errors.GenericError("Unknown command %r", cmd)
- def _StartSession(self, args):
- """Starts a new session.
- @type args: dict
- @param args: Arguments passed to command by client
- """
- return self._StartSessionInner(args, None)
- def _AttachSession(self, args, shadowcookie):
- """Attaches to an existing session, shadowing it.
- @type args: dict
- @param args: Arguments passed to command by client
- @type shadowcookie: str
- @param shadowcookie: Session cookie for session to be shadowed
- """
- assert shadowcookie
- logging.debug("Attaching to session with shadowcookie %r", shadowcookie)
- return self._StartSessionInner(args, shadowcookie)
- def _StartSessionInner(self, args, shadowcookie):
- ctx = self._ctx
- if ctx.sessrunner:
- raise errors.GenericError("Session already started")
- ctx.session = node.NodeSession(ctx, args)
- if shadowcookie:
- ctx.session.SetShadowCookie(shadowcookie)
- sessrunner = node.SessionRunner(ctx)
- sessrunner.Start()
- ctx.sessrunner = sessrunner
- return True
- def _RestoreSession(self, args):
- """Restores a session.
- @type args: dict
- @param args: Arguments passed to command by client
- """
- ctx = self._ctx
- if not ctx.sessrunner:
- raise errors.GenericError("Session not yet started")
- logging.debug("Resuming session")
- ctx.session.PrepareRestore(args)
- ctx.sessrunner.Restore()
- return True
- def _TerminateSession(self, args):
- """Terminates the current session.
- """
- raise NotImplementedError()
- def _GetShadowCookie(self):
- """Returns the cookie needed to shadow the current session.
- @rtype: str
- @return: Shadow cookie
- """
- ctx = self._ctx
- if not ctx.sessrunner:
- raise errors.GenericError("Session not yet started")
- # TODO: If request is coming from a different user, show dialog before
- # giving out cookie. This is not a problem at the moment--only the user
- # themselves can access the node socket.
- return ctx.session.cookie
- class ClientConnection:
- def __init__(self, ctx):
- self._ops = ClientOperations(ctx)
- self.__conn = None
- self.__channel = daemon.IOChannel()
- self.__reader = daemon.ChopReader(node.PROTO_SEPARATOR)
- signal_name = daemon.ChopReader.SLICE_COMPLETE_SIGNAL
- self.__reader_slice_reg = \
- daemon.SignalRegistration(self.__reader,
- self.__reader.connect(signal_name,
- self.__HandleSlice))
- self.__reader.Attach(self.__channel)
- def Attach(self, conn):
- self.__conn = conn
- self.__channel.Attach(conn.fileno())
- def __del__(self):
- # TODO: Close
- pass
- def __HandleSlice(self, _, data):
- success = False
- try:
- req = serializer.LoadJson(data)
- ValidateRequest(req)
- cmd = req[node.REQ_FIELD_CMD]
- args = req[node.REQ_FIELD_ARGS]
- # Call function
- result = self._ops(cmd, args)
- success = True
- except (SystemExit, KeyboardInterrupt):
- raise
- except errors.GenericError, err:
- # Serialize exception arguments
- result = (err.__class__.__name__, err.args)
- except Exception, err:
- logging.exception("Error while handling request")
- result = "Caught exception: %s" % str(err)
- response = {
- node.RESP_FIELD_SUCCESS: success,
- node.RESP_FIELD_RESULT: result,
- }
- serialized_data = serializer.DumpJson(response)
- assert node.PROTO_SEPARATOR not in serialized_data
- self.__channel.Write(serialized_data + node.PROTO_SEPARATOR)
- class NodeSocket:
- def __init__(self, ctx, path):
- self.__ctx = ctx
- self.__path = path
- self.__socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
- def Start(self):
- self.__socket.bind(self.__path)
- self.__socket.listen(32)
- gobject.io_add_watch(self.__socket.fileno(), gobject.IO_IN,
- self.__HandleIO)
- def __HandleIO(self, source, cond):
- if cond & gobject.IO_IN:
- self.__IncomingConnection()
- return True
- return False
- def __IncomingConnection(self):
- (conn, _) = self.__socket.accept()
- logging.info("Connection established")
- ClientConnection(self.__ctx).Attach(conn)
- def _CheckIfSessionWasStarted(ctx):
- if not ctx.sessrunner:
- logging.error("Session wasn't started in %s seconds, terminating",
- _SESSION_START_TIMEOUT)
- sys.exit(1)
- return False
- class NxNodeProgram(cli.GenericProgram):
- def Run(self):
- if len(self.args) != 2:
- raise errors.GenericError("Username or session ID missing")
- ctx = NxNodeContext()
- ctx.cfg = self.cfg
- ctx.sessmgr = session.NxSessionManager()
- ctx.processes = []
- (ctx.username, ctx.sessid) = self.args
- ctx.uid = _GetUserUid(ctx.username)
- server = NodeSocket(ctx, ctx.sessmgr.GetSessionNodeSocket(ctx.sessid))
- server.Start()
- # Terminate if session wasn't started after some time
- gobject.timeout_add(_SESSION_START_TIMEOUT * 1000,
- _CheckIfSessionWasStarted, ctx)
- mainloop = gobject.MainLoop()
- logging.debug("Starting mainloop")
- mainloop.run()
- def Main():
- logsetup = utils.LoggingSetup(PROGRAM)
- NxNodeProgram(logsetup).Main()