/neatx/lib/app/nxnode.py

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