cwmp

#!/usr/bin/env python
# Copyright 2011 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Command-line client for tr/rcommand.py."""

__author__ = 'apenwarr@google.com (Avery Pennarun)'

import os.path
import re
import socket
import sys
import traceback
import google3
import bup.options
import bup.shquote
import tr.mainloop
import tr.quotedblock

try:
  import readline  # pylint:disable=g-import-not-at-top
except ImportError:
  readline = None


optspec = """
cwmp [options] [command...]
--
u,unix-path=  Unix socket server is listening on [/tmp/cwmpd.sock]
i,ip=        IP hostname/ipaddr server is listening on (default: unix socket)
p,port=      IP port server is listening on [12999]
"""


HISTORY_FILE = os.path.expanduser('/tmp/.cwmp_history')


def Log(s, *args):
  s = str(s)
  if args:
    sys.stderr.write((s + '\n') % args)
  else:
    sys.stderr.write(s + '\n')


_want_verbose = False


def Verbose(s, *args):
  if _want_verbose:
    Log(s, *args)


class Fatal(Exception):
  pass


def HandleFatal(func):
  def Fn(*args, **kwargs):
    try:
      return func(*args, **kwargs)
    except Fatal, e:
      Log('Fatal: %s' % e)
      sys.exit(1)
  return Fn


def _NormalizePath(path):
  """Like os.path.normpath, but doesn't remove any trailing slash."""
  result = os.path.normpath(path)
  if path.endswith('/') and not result.endswith('/'):
    result += '/'
  return result


def _DotsToSlashes(s):
  return re.sub(r'([^/.])\.', r'\1/', s)


def _SlashesToDots(s):
  name = s.replace('/', '.')
  if name.startswith('.'):
    name = name[1:]
  return name


class Client(object):
  """Manage the client-side state of an rcommand connection."""

  def __init__(self, loop, connector):
    self.loop = loop
    self.connector = connector
    self.stream = None
    self.result = None
    self._last_res = (None, None, None)
    self.cwd = '/'
    self.quotedblock = tr.quotedblock.QuotedBlockProtocol(
        HandleFatal(self.GotBlock))
    self._StartConnect()

  def _StartConnect(self):
    self.stream = None
    try:
      self.connector(HandleFatal(self.OnConnect))
    except socket.error, e:
      raise Fatal(str(e))

  def Close(self):
    if self.stream:
      self.stream.close()

  def OnConnect(self, stream):
    if not stream:
      raise Fatal('connection failed')
    Verbose('Connected to server.\n')
    self.stream = stream
    self.stream.set_close_callback(HandleFatal(self.OnClose))
    self._StartRead()
    self.loop.ioloop.stop()

  def _StartRead(self):
    self.stream.read_until('\r\n\r\n', HandleFatal(self.GotData))

  def OnClose(self):
    Log('Server connection closed!')
    self._StartConnect()

  def GotData(self, data):
    assert data.endswith('\r\n\r\n')
    for i in data.split('\r\n')[:-1]:
      self.quotedblock.GotData(i)
    self._StartRead()

  def GotBlock(self, lines):
    self.result = lines
    self.loop.ioloop.stop()

  def Send(self, lines):
    s = self.quotedblock.RenderBlock(lines)
    self.stream.write(s)

  def Run(self, lines):
    self.Send(lines)
    self.loop.Start()
    result = self.result
    self.result = None
    return result

  def _RequestCompletions(self, prefix):
    prefix = _NormalizePath(prefix)
    completions = self.Run([['completions', _SlashesToDots(prefix)]])[1:]
    for [i] in completions:
      yield i

  def _GetSubstitutions(self, line):
    (qtype, lastword) = bup.shquote.unfinished_word(line)
    request = os.path.join(self.cwd, _DotsToSlashes(lastword))
    subs = list(self._RequestCompletions(request))
    cmd = line.split(' ', 1)[0]
    if cmd.lower() in ('cd', 'ls', 'list', 'rlist', 'validate', 'add', 'del'):
      # only return object names, not parameters
      subs = [i for i in subs if i.endswith('.')]
    return (qtype, lastword, subs)

  def _StripPathPrefix(self, oldword, newword):
    # readline is weird: we have to remove all the parts that were before
    # the last '/', but not parts before the last '.', because we have to
    # tell it what to replace everything after the last '/' with.
    after_slash = oldword.split('/')[-1]
    dots = after_slash.split('.')
    if newword.endswith('.'):
      new_last_dot = '.'.join(newword.split('.')[-2:])
    else:
      new_last_dot = newword.split('.')[-1]
    dots[-1] = new_last_dot
    return '.'.join(dots)

  def ReadlineCompleter(self, text, state):
    """Callback for the readline library to autocomplete a line of text.

    Args:
      text: the current input word (basename following the last slash)
      state: a number of 0..n, where n is the number of substitutions.
    Returns:
      One of the available substitutions.
    """
    try:
      text = _DotsToSlashes(text)
      line = readline.get_line_buffer()[:readline.get_endidx()]
      if not state:
        self._last_res = self._GetSubstitutions(line)
      (qtype, lastword, subs) = self._last_res
      if state < len(subs):
        new_last_slash = _DotsToSlashes(self._StripPathPrefix(lastword,
                                                              subs[state]))
        is_param = not new_last_slash.endswith('/')
        if is_param and qtype:
          new_last_slash += qtype
        return new_last_slash
    except Exception, e:  # pylint:disable=broad-except
      Log('\n')
      try:
        traceback.print_tb(sys.exc_traceback)
      except Exception, e2:  # pylint:disable=broad-except
        Log('Error printing traceback: %s\n' % e2)
    Log('\nError in completion: %s\n' % e)


def DoCmd(client, words):
  """Send the command to cwmpd."""
  words = [i.decode('utf-8') for i in words]
  cmd, args = (words[0].lower(), words[1:])
  if cmd in ('cd', 'ls', 'list', 'rlist', 'validate',
             'add', 'del', 'get', 'set'):
    if not args:
      args = [client.cwd]
    skip = 2 if cmd in ('add', 'del', 'set') else 1
    for i in range(0, len(args), skip):
      relpath = _DotsToSlashes(args[i])
      abspath = os.path.normpath(os.path.join(client.cwd, relpath))
      args[i] = _SlashesToDots(abspath)
  if cmd == 'cd':
    client.cwd = os.path.normpath(os.path.join(client.cwd, relpath))
  else:
    line = [cmd] + args
    result = client.Run([line])
    return result


def Interactive(client):
  """Run an interactive command prompt on the given connection."""
  global _want_verbose
  _want_verbose = True
  if readline:
    readline.set_completer_delims(' \t\n\r/')
    readline.set_completer(client.ReadlineCompleter)
    # this. drove. me. INSANE.
    # http://stackoverflow.com/questions/7124035
    if 'libedit' in readline.__doc__:
      readline.parse_and_bind('bind ^I rl_complete')  # MacOS
    else:
      readline.parse_and_bind('tab: complete')        # other

  while True:
    print
    line = raw_input('%s> ' % client.cwd) + '\n'
    while 1:
      word = bup.shquote.unfinished_word(line)[1]
      if not word:
        break
      line += raw_input('%*s> ' % (len(client.cwd), '')) + '\n'
    words = [word for (unused_idx, word) in bup.shquote.quotesplit(line)]
    if not words:
      continue
    result = DoCmd(client, words)
    print client.quotedblock.RenderBlock(result).strip()


def main():
  o = bup.options.Options(optspec)
  (opt, unused_flags, extra) = o.parse(sys.argv[1:])

  if readline and os.path.exists(HISTORY_FILE):
    readline.read_history_file(HISTORY_FILE)
  client = None
  try:
    loop = tr.mainloop.MainLoop()
    if opt.ip is not None:
      connector = lambda err: loop.ConnectInet((opt.ip, opt.port), err)
    else:
      connector = lambda err: loop.ConnectUnix(opt.unix_path, err)
    client = Client(loop, connector)
    loop.Start()
    if extra:
      result = DoCmd(client, extra)
      code = result.pop(0)
      if code[0] != 'OK':
        raise Fatal(' '.join(code))
      for line in result:
        print ' '.join(line)
    else:
      Interactive(client)
  except Fatal, e:
    Log(e)
    sys.exit(1)
  except EOFError:
    pass
  finally:
    if readline:
      readline.write_history_file(HISTORY_FILE)
    if client:
      client.Close()


if __name__ == '__main__':
  main()