/neatx/lib/protocol.py

  1#
  2#
  3
  4# Copyright (C) 2009 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"""NX protocol utilities
 23
 24"""
 25
 26
 27import logging
 28import re
 29import urllib
 30
 31from neatx import utils
 32
 33
 34NX_PROMPT = "NX>"
 35NX_EOL = "\n"
 36NX_EOL_CHARS = NX_EOL
 37
 38NX_CMD_HELLO = "hello"
 39NX_CMD_BYE = "bye"
 40NX_CMD_LOGIN = "login"
 41NX_CMD_LISTSESSION = "listsession"
 42NX_CMD_STARTSESSION = "startsession"
 43NX_CMD_ATTACHSESSION = "attachsession"
 44NX_CMD_RESTORESESSION = "restoresession"
 45NX_CMD_TERMINATE = "terminate"
 46NX_CMD_SET = "set"
 47NX_CMD_QUIT = "quit"
 48
 49NX_FALSE = "0"
 50NX_TRUE = "1"
 51
 52
 53class NxQuitServer(Exception):
 54  pass
 55
 56
 57class NxQuietQuitServer(NxQuitServer):
 58  pass
 59
 60
 61class NxProtocolError(Exception):
 62  def __init__(self, code, message, fatal=False):
 63    self.code = code
 64    self.msg = message
 65    self.fatal = fatal
 66
 67
 68class NxUndefinedCommand(NxProtocolError):
 69  def __init__(self, command):
 70    NxProtocolError.__init__(self, 503,
 71                             "Error: undefined command: '%s'" % command)
 72
 73
 74class NxNotBeforeLogin(NxProtocolError):
 75  def __init__(self, command):
 76    message = ("Error: the command '%s' cannot be called before to login" %
 77               command)
 78    NxProtocolError.__init__(self, 554, message)
 79
 80
 81class NxNotAfterLogin(NxProtocolError):
 82  def __init__(self, command):
 83    message = "Error: the command '%s' cannot be called after login" % command
 84    NxProtocolError.__init__(self, 554, message)
 85
 86
 87class NxUnsupportedProtocol(NxProtocolError):
 88  def __init__(self):
 89    # Had to set code to 500 instead of 552, otherwise client ignores
 90    # this error.
 91    NxProtocolError.__init__(self, 500,
 92                             ("Protocol you requested is not supported, "
 93                              "please upgrade your client to latest version"),
 94                             True)
 95
 96
 97class NxUnencryptedSessionsNotAllowed(NxProtocolError):
 98  def __init__(self, x):
 99    message = "ERROR: Unencrypted sessions are not allowed on this server"
100    NxProtocolError.__init__(self, 594, message)
101
102
103class NxParameterParsingError(NxProtocolError):
104  def __init__(self, params):
105    message = (("Error: Parsing parameters: string \"%s\" has "
106                "invalid format") % params)
107    NxProtocolError.__init__(self, 597, message)
108
109
110class NxServerBase(object):
111  """Base class for NX protocol servers.
112
113  """
114  def __init__(self, input, output, handler):
115    """Instance initialization.
116
117    @type input: file
118    @param input: Input file handle
119    @type output: file
120    @param output: Output file handle
121    @type handler: callable
122    @param handler: Called for received lines
123
124    """
125    assert callable(handler)
126
127    self._input = input
128    self._output = output
129    self._handler = handler
130
131  def Start(self):
132    """Start responding to requests.
133
134    """
135    self.SendBanner()
136
137    while True:
138      self.Write(105)
139      try:
140        try:
141          line = self.ReadLine()
142
143          # Ignore empty lines
144          if line.strip():
145            self._HandleLine(line)
146
147        except NxProtocolError, err:
148          self.Write(err.code, message=err.msg)
149          if err.fatal:
150            raise NxQuitServer()
151
152      except NxQuietQuitServer:
153        break
154
155      except NxQuitServer:
156        self.Write(999, "Bye.")
157        break
158
159  def _HandleLine(self, line):
160    try:
161      self._handler(line)
162    except (SystemExit, KeyboardInterrupt, NxProtocolError, NxQuitServer):
163      raise
164    except Exception:
165      logging.exception("Error while handling line %r", line)
166      raise NxProtocolError(500, "Internal error", fatal=True)
167
168  def _Write(self, data):
169    """Write to output after logging.
170
171    """
172    logging.debug(">>> %r", data)
173    try:
174      self._output.write(data)
175    finally:
176      self._output.flush()
177
178  def Write(self, code, message=None, newline=None):
179    """Write prompt to output.
180
181    @type code: int
182    @param code: Status code
183    @type message: str
184    @param message: Message text
185    @type newline: bool
186    @param newline: Whether to add newline
187
188    Note: The "newline" parameter is a tri-state variable. If there's a
189    message, print newline by default (e.g. "NX> 500 something\\n"). If there's
190    no message, don't print newline (e.g. "NX> 105 "). This logic can be
191    overridden by explictly setting the "newline" parameter to a non-None
192    value.
193
194    """
195    assert code >= 0 and code <= 999
196
197    # Build prompt
198    prompt = "%s %s " % (NX_PROMPT, code)
199    if message:
200      prompt += "%s" % message
201    if (newline is None and message) or (newline is not None and newline):
202      prompt += NX_EOL
203
204    self._Write(prompt)
205
206  def WriteLine(self, line):
207    """Write line to output.
208
209    One newline char is automatically added.
210
211    """
212    self._Write(line + NX_EOL)
213
214  def ReadLine(self, hide=False):
215    """Reads line from input.
216
217    @type hide: bool
218    @param hide: Whether to hide line read from log output
219
220    """
221    # TODO: Timeout (poll, etc.)
222    line = self._input.readline()
223
224    if hide:
225      logging.debug("<<< [hidden]")
226    else:
227      logging.debug("<<< %r", line)
228
229    # Has the client closed the connection?
230    if not line:
231      raise NxQuitServer()
232
233    return line.rstrip(NX_EOL_CHARS)
234
235  def WithoutTerminalEcho(self, fn, *args, **kwargs):
236    """Calls function with ECHO flag disabled.
237
238    @type fn: callable
239    @param fn: Called function
240
241    """
242    return utils.WithoutTerminalEcho(self._input, fn, *args, **kwargs)
243
244  def SendBanner(self):
245    """Send banner to peer.
246
247    Can be overriden by subclass.
248
249    """
250
251
252def SplitCommand(command):
253  """Split line into command and arguments on first whitespace.
254
255  """
256  parts = command.split(None, 1)
257
258  # Empty lines should've been filtered out earlier
259  assert parts
260
261  if len(parts) == 1:
262    args = ""
263  else:
264    args = parts[1]
265
266  return (parts[0].lower(), args)
267
268
269def ParseParameters(params, _logging=logging):
270  """Parse parameters sent by client.
271
272  @type params: string
273  @param params: Parameter string
274
275  """
276  param_re = re.compile((r"^\s*--(?P<name>[a-z][a-z0-9_-]*)="
277                         "\"(?P<value>[^\"]*)\"\\s*"), re.I)
278  result = []
279  work = params.strip()
280
281  while work:
282    m = param_re.search(work)
283    if not m:
284      _logging.warning("Failed to parse parameter string %r", params)
285      raise NxParameterParsingError(params)
286
287    result.append((m.group("name"), m.group("value")))
288    work = work[m.end():]
289
290  assert not work
291
292  return result
293
294
295def UnquoteParameterValue(value):
296  """Unquotes parameter value.
297
298  @type value: str
299  @param value: Quoted value
300  @rtype: str
301  @return: Unquoted value
302
303  """
304  return urllib.unquote(value)
305
306
307def ParseNxBoolean(value):
308  """Parses a boolean parameter value.
309
310  @type value: str
311  @param value: Value
312  @rtype: bool
313  @return: Whether parameter evaluates to true
314
315  """
316  return value == NX_TRUE
317
318
319def FormatNxBoolean(value):
320  """Format boolean value for nxagent.
321
322  @type value: bool
323  @param value: Value
324  @rtype: str
325
326  """
327  if value:
328    return NX_TRUE
329
330  return NX_FALSE
331
332
333def ParseNxSize(value):
334  """Parses a size unit parameter value.
335
336  @type value: str
337  @param value: Value
338  @rtype: int
339  @return: Size in Mebibytes
340
341  """
342  return int(value.rstrip("M"))
343
344
345def FormatNxSize(value):
346  """Format size value.
347
348  @type value: int
349  @param value: Value in Mebibytes
350  @rtype: str
351
352  """
353  return "%dM" % value