/neatx/lib/app/nxnode.py

http://neatx.googlecode.com/ · Python · 307 lines · 230 code · 17 blank · 60 comment · 3 complexity · cee8c2a33a01a67805d5d42d25d5665f MD5 · raw file

  1. #
  2. #
  3. # Copyright (C) 2007 Google Inc.
  4. #
  5. # This program is free software; you can redistribute it and/or modify
  6. # it under the terms of the GNU General Public License as published by
  7. # the Free Software Foundation; either version 2 of the License, or
  8. # (at your option) any later version.
  9. #
  10. # This program is distributed in the hope that it will be useful, but
  11. # WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  13. # General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU General Public License
  16. # along with this program; if not, write to the Free Software
  17. # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
  18. # 02110-1301, USA.
  19. """nxnode program.
  20. This program is started once per session. It receives commands from nxserver
  21. via a Unix socket and updates the session database based on nxagent's output.
  22. """
  23. import logging
  24. import pwd
  25. import select
  26. import signal
  27. import socket
  28. import sys
  29. import gobject
  30. from neatx import cli
  31. from neatx import constants
  32. from neatx import daemon
  33. from neatx import errors
  34. from neatx import node
  35. from neatx import serializer
  36. from neatx import session
  37. from neatx import utils
  38. PROGRAM = "nxnode"
  39. _SESSION_START_TIMEOUT = 30
  40. class NxNodeContext(object):
  41. def __init__(self):
  42. self.cfg = None
  43. self.uid = None
  44. self.username = None
  45. self.sessrunner = None
  46. self.sessid = None
  47. self.session = None
  48. self.sessmgr = None
  49. self.processes = None
  50. def _GetUserUid(username):
  51. return pwd.getpwnam(username).pw_uid
  52. def ValidateRequest(req):
  53. if not (isinstance(req, dict) and
  54. node.REQ_FIELD_CMD in req and
  55. node.REQ_FIELD_ARGS in req):
  56. raise errors.GenericError("Incomplete request")
  57. class ClientOperations(object):
  58. def __init__(self, ctx):
  59. self._ctx = ctx
  60. def __call__(self, cmd, args):
  61. logging.info("Received request: %r, %r", cmd, args)
  62. if cmd == node.CMD_STARTSESSION:
  63. return self._StartSession(args)
  64. elif cmd == node.CMD_ATTACHSESSION:
  65. assert len(args) == 2
  66. return self._AttachSession(args[0], args[1])
  67. elif cmd == node.CMD_RESTORESESSION:
  68. return self._RestoreSession(args)
  69. elif cmd == node.CMD_TERMINATESESSION:
  70. return self._TerminateSession(args)
  71. elif cmd == node.CMD_GET_SHADOW_COOKIE:
  72. return self._GetShadowCookie()
  73. else:
  74. raise errors.GenericError("Unknown command %r", cmd)
  75. def _StartSession(self, args):
  76. """Starts a new session.
  77. @type args: dict
  78. @param args: Arguments passed to command by client
  79. """
  80. return self._StartSessionInner(args, None)
  81. def _AttachSession(self, args, shadowcookie):
  82. """Attaches to an existing session, shadowing it.
  83. @type args: dict
  84. @param args: Arguments passed to command by client
  85. @type shadowcookie: str
  86. @param shadowcookie: Session cookie for session to be shadowed
  87. """
  88. assert shadowcookie
  89. logging.debug("Attaching to session with shadowcookie %r", shadowcookie)
  90. return self._StartSessionInner(args, shadowcookie)
  91. def _StartSessionInner(self, args, shadowcookie):
  92. ctx = self._ctx
  93. if ctx.sessrunner:
  94. raise errors.GenericError("Session already started")
  95. ctx.session = node.NodeSession(ctx, args)
  96. if shadowcookie:
  97. ctx.session.SetShadowCookie(shadowcookie)
  98. sessrunner = node.SessionRunner(ctx)
  99. sessrunner.Start()
  100. ctx.sessrunner = sessrunner
  101. return True
  102. def _RestoreSession(self, args):
  103. """Restores a session.
  104. @type args: dict
  105. @param args: Arguments passed to command by client
  106. """
  107. ctx = self._ctx
  108. if not ctx.sessrunner:
  109. raise errors.GenericError("Session not yet started")
  110. logging.debug("Resuming session")
  111. ctx.session.PrepareRestore(args)
  112. ctx.sessrunner.Restore()
  113. return True
  114. def _TerminateSession(self, args):
  115. """Terminates the current session.
  116. """
  117. raise NotImplementedError()
  118. def _GetShadowCookie(self):
  119. """Returns the cookie needed to shadow the current session.
  120. @rtype: str
  121. @return: Shadow cookie
  122. """
  123. ctx = self._ctx
  124. if not ctx.sessrunner:
  125. raise errors.GenericError("Session not yet started")
  126. # TODO: If request is coming from a different user, show dialog before
  127. # giving out cookie. This is not a problem at the moment--only the user
  128. # themselves can access the node socket.
  129. return ctx.session.cookie
  130. class ClientConnection:
  131. def __init__(self, ctx):
  132. self._ops = ClientOperations(ctx)
  133. self.__conn = None
  134. self.__channel = daemon.IOChannel()
  135. self.__reader = daemon.ChopReader(node.PROTO_SEPARATOR)
  136. signal_name = daemon.ChopReader.SLICE_COMPLETE_SIGNAL
  137. self.__reader_slice_reg = \
  138. daemon.SignalRegistration(self.__reader,
  139. self.__reader.connect(signal_name,
  140. self.__HandleSlice))
  141. self.__reader.Attach(self.__channel)
  142. def Attach(self, conn):
  143. self.__conn = conn
  144. self.__channel.Attach(conn.fileno())
  145. def __del__(self):
  146. # TODO: Close
  147. pass
  148. def __HandleSlice(self, _, data):
  149. success = False
  150. try:
  151. req = serializer.LoadJson(data)
  152. ValidateRequest(req)
  153. cmd = req[node.REQ_FIELD_CMD]
  154. args = req[node.REQ_FIELD_ARGS]
  155. # Call function
  156. result = self._ops(cmd, args)
  157. success = True
  158. except (SystemExit, KeyboardInterrupt):
  159. raise
  160. except errors.GenericError, err:
  161. # Serialize exception arguments
  162. result = (err.__class__.__name__, err.args)
  163. except Exception, err:
  164. logging.exception("Error while handling request")
  165. result = "Caught exception: %s" % str(err)
  166. response = {
  167. node.RESP_FIELD_SUCCESS: success,
  168. node.RESP_FIELD_RESULT: result,
  169. }
  170. serialized_data = serializer.DumpJson(response)
  171. assert node.PROTO_SEPARATOR not in serialized_data
  172. self.__channel.Write(serialized_data + node.PROTO_SEPARATOR)
  173. class NodeSocket:
  174. def __init__(self, ctx, path):
  175. self.__ctx = ctx
  176. self.__path = path
  177. self.__socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  178. def Start(self):
  179. self.__socket.bind(self.__path)
  180. self.__socket.listen(32)
  181. gobject.io_add_watch(self.__socket.fileno(), gobject.IO_IN,
  182. self.__HandleIO)
  183. def __HandleIO(self, source, cond):
  184. if cond & gobject.IO_IN:
  185. self.__IncomingConnection()
  186. return True
  187. return False
  188. def __IncomingConnection(self):
  189. (conn, _) = self.__socket.accept()
  190. logging.info("Connection established")
  191. ClientConnection(self.__ctx).Attach(conn)
  192. def _CheckIfSessionWasStarted(ctx):
  193. if not ctx.sessrunner:
  194. logging.error("Session wasn't started in %s seconds, terminating",
  195. _SESSION_START_TIMEOUT)
  196. sys.exit(1)
  197. return False
  198. class NxNodeProgram(cli.GenericProgram):
  199. def Run(self):
  200. if len(self.args) != 2:
  201. raise errors.GenericError("Username or session ID missing")
  202. ctx = NxNodeContext()
  203. ctx.cfg = self.cfg
  204. ctx.sessmgr = session.NxSessionManager()
  205. ctx.processes = []
  206. (ctx.username, ctx.sessid) = self.args
  207. ctx.uid = _GetUserUid(ctx.username)
  208. server = NodeSocket(ctx, ctx.sessmgr.GetSessionNodeSocket(ctx.sessid))
  209. server.Start()
  210. # Terminate if session wasn't started after some time
  211. gobject.timeout_add(_SESSION_START_TIMEOUT * 1000,
  212. _CheckIfSessionWasStarted, ctx)
  213. mainloop = gobject.MainLoop()
  214. logging.debug("Starting mainloop")
  215. mainloop.run()
  216. def Main():
  217. logsetup = utils.LoggingSetup(PROGRAM)
  218. NxNodeProgram(logsetup).Main()