Ticket #45: yamlrpclib.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
# A port of xmlrpclib to jsonrpclib to yamlrpclib....
#
#
# The YAML-RPC client interface is based on the JSON-RPC client
# which is based on the XML-RPC client.
#
# Copyright (c) 1999-2002 by Secret Labs AB
# Copyright (c) 1999-2002 by Fredrik Lundh
# Copyright (c) 2006 by Matt Harrison
# Copyright (c) 2007 by Peter Murphy
#
# By obtaining, using, and/or copying this software and/or its
# associated documentation, you agree that you have read, understood,
# and will comply with the following terms and conditions:
#
# Permission to use, copy, modify, and distribute this software and
# its associated documentation for any purpose and without fee is
# hereby granted, provided that the above copyright notice appears in
# all copies, and that both that copyright notice and this permission
# notice appear in supporting documentation, and that the name of
# Secret Labs AB or the author not be used in advertising or publicity
# pertaining to distribution of the software without specific, written
# prior permission.
#
# SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD
# TO THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANT-
# ABILITY AND FITNESS.  IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR
# BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY
# DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
# WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
# ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
# OF THIS SOFTWARE.
# --------------------------------------------------------------------

import urllib
import httplib
import base64
import yaml
import yamlrpcbasic
import types
from xmlrpclib import _Method;
from yamlrpcbasic import getAttrWay;

# Implementation version.

__version__ = "0.0.1"

# Protocol implementation.

protvers = "0.1"

# We add error handling here.

class YRPCException(Exception):
    """ Base class for all YAML-RPC errors."""
    def __str__(self):
        return repr(self);
    def __unicode__(self):
        return repr(self);

class YRPCProtocolException(YRPCException):
    """ Indicates an HTTP/HTTPS protocol error. This is raised by the HTTP
        transport layer, if the server returns an error code other than 200
        (OK). The parameters are:

        url: the target URL.
        errcode: The HTTP or HTTPS error code.
        errmsg: The HTTP or HTTPS error message.
        headers The HTTP or HTTPS header dictionary.
    """
    def __init__(self, url, errcode, errmsg, headers):
        YRPCException.__init__(self)
        self.url = url
        self.errcode = errcode
        self.errmsg = errmsg
        self.headers = headers
    def __repr__(self):
        return ("<YAML-RPC ProtocolError for %s: %s %s>" %
            (self.url, self.errcode, self.errmsg))

class YRPCErrorException(YRPCException):
    """ This represents an YAML-RPC error object. This object is thrown when
        the YAML-RPC server sends an error object back as YAML-RPC data.
        For example, a called routine may not exist on the server, or there
        are not enough parameters for the routine.
    """
    def __init__(self, errorObj):
        """ One parameter. This is either an object of type YAMLRPCError, or
            a dictionary with "name", "code", "message" and "error" members.
        """
        YRPCException.__init__(self);
        self.errorObj = errorObj;
    def __repr__(self):
        if not self.errorObj:
            return None;
        obj = self.errorObj;
        return ("<YAML-RPC %s (%s): %s>" %
            (getAttrWay(obj, "code"), getAttrWay(obj, "name"), 
            getAttrWay(obj, "message")))

ID = 1
def _gen_id():
    """ Generates an ID number for each YAML-RPC procedure call. """ 
    global ID
    ID = ID + 1
    return ID

class YAMLTransport:
    """ This handles an HTTP transaction to an YAML-RPC server."""

    # The client identifier (may be overridden).
    user_agent = "yamlrpclib.py/%s (by Peter Murphy)" % __version__

    def __init__(self, safe_loader = True):
        """ The initializer. There is one parameter for this function:
            
            safe_loader: the default is True: only allows "simple" YAML data to
            be loaded - sequences, dictionaries and most immutable types. If
            this parameter is False, then arbitrary YAML data can be loaded.
            This can be a security risk. 
            
            *** Only set this parameter to False with trusted servers! ***
        """
        self.loadfunc(safe_loader);

    def loadfunc(self, safe_loader):
        """ Sets the load function. """
        self.safe_loader = safe_loader;
        if self.safe_loader:
            self._loadfunc = yaml.safe_load;
        else:
            self._loadfunc = yaml.load;

    def request(self, host, handler, request_body, verbose=0):
        """ This sends a complete request, and parses the response. The 
            parameters are:
            
            host: the trget host.
            handler: the target RPC handler.
            request_body: a YAML-RPC request body.
            verbose: the debugging flag.
            
            The routine returns the parsed response.
        """
        h = self.make_connection(host)
        if verbose:
            h.set_debuglevel(1)

        self.send_request(h, handler, request_body)
        self.send_host(h, host)
        self.send_user_agent(h)
        self.send_content(h, request_body)

        errcode, errmsg, headers = h.getreply()

        if errcode != 200:
            raise YRPCProtocolException(host + handler, errcode, errmsg,
                headers);
        self.verbose = verbose

        try:
            sock = h._conn.sock
        except AttributeError:
            sock = None

        return self._parse_response(h.getfile(), sock)

    def get_host_info(self, host):
        """ This gets authorization info from the host parameter. This may
            be a string, or a (host, x509-dict) tuple. If it is a string,
            it is checked for a "user:pw@host" format, and a "Basic
            Authentication" header is added if appropriate.
    
            The function returns a 3-tuple of the form:
            
            (actual host, extra headers, x509 info).  
            
            The header and x509 fields may be None.
        """

        x509 = {}
        if isinstance(host, types.TupleType):
            host, x509 = host

        auth, host = urllib.splituser(host)

        if auth:
            auth = base64.encodestring(urllib.unquote(auth))
            auth = string.join(string.split(auth), "") 
            # Done to get rid of whitespace
            extra_headers = [("Authorization", "Basic " + auth)]
        else:
            extra_headers = None

        return host, extra_headers, x509

    def make_connection(self, host):
        """ This creates a HTTP connection object from a host descriptor.
        
            We know "host" is the target host. This returns a connection
            handle.
        """
        host, extra_headers, x509 = self.get_host_info(host)
        return httplib.HTTP(host)

    def send_request(self, connection, handler, request_body):
        """ This sends a request header. The parameters are:
        
            connection: the connection handle.
            handler: the target RPC handler.
            request_body: the YAML-RPC body.
        """
        connection.putrequest("POST", handler)

    def send_host(self, connection, host):
        """ Sends by host name. Parameters:
        
            connection: connection handle.
            host: host name.
        """
        host, extra_headers, x509 = self.get_host_info(host)
        connection.putheader("Host", host)
        if extra_headers:
            if isinstance(extra_headers, DictType):
                extra_headers = extra_headers.items()
            for key, value in extra_headers:
                connection.putheader(key, value)

    def send_user_agent(self, connection):
        """ This sends the user-agent identifier. The sole parameter is 
            connection: the connection handle.
        """
        connection.putheader("User-Agent", self.user_agent)

    def send_content(self, connection, request_body):
        """ This sends the request body. There are two parameters:
        
            connection: the connection handle.
            request_body: the YAML-RPC request body.
        """
        connection.putheader("Content-Type", "text/yaml")
        connection.putheader("Content-Length", str(len(request_body)))
        connection.endheaders()
        if request_body:
            connection.send(request_body)

    def parse_response(self, file):
        """ This parses the response from a YAML-RPC server. The sole parameter
            is file (a stream representing the input socket). 
            
            This function returns the YAML data loaded as an object or a dict.

            Note: this is provided as a compatibility interface.
        """
        return self._parse_response(file, None)

    def _parse_response(self, file, sock):
        """ This is the alternative interface for parsing a response. It is
            like the parse_response method, but also provides direct access to
            the underlying socket object (where available). Parameters:
            
            file: a socket as a file stream.
            
            sock: the socket handle (or None, if the socket object could not
            be accessed).

            This function returns the YAML data loaded as an object or a dict.
        """
        fedin = "";
        while 1:
            if sock:
                response = sock.recv(1024);
            else:
                response = file.read(1024);
            if not response:
                break;
            if self.verbose:
                print "body: ", repr(response);
            fedin = fedin + response;
        file.close();
        return self._loadfunc(fedin);


class SecureYAMLTransport(YAMLTransport):
    """ This handles an HTTPS transaction to an YAML-RPC server.
    
        Warning: this is untested! The code is based on SafeTransport from
        jsonrpclib.py.
    """

    def __init__(self, safe_loader = True):
        """ The initializer. There is one parameter for this function:
            
            safe_loader: the default is True: only allows "simple" YAML data to
            be loaded - sequences, dictionaries and most immutable types. If
            this parameter is False, then arbitrary YAML data can be loaded.
            This can be a security risk. 
            
            *** Only set this parameter to False with trusted servers! ***
        """
        YAMLTransport.__init__(self, safe_loader);

    def make_connection(self, host):
        """ This creates a HTTPS connection object from a host descriptor.
        
            host may be a string, or a (host, x509-dict) tuple.
        """
        host, extra_headers, x509 = self.get_host_info(host)
        try:
            HTTPS = httplib.HTTPS
        except AttributeError:
            raise NotImplementedError(
                "your version of httplib doesn't support HTTPS"
                )
        else:
            return HTTPS(host, None, **(x509 or {}))


class YAMLServerProxy(object):
    """ This acts as a proxy for calling YAML-RPC methods on a remote server,
        in the same way that ServerProxy allows one to access XML-RPC methods
        on a remote server. For example:

        s = YAMLServerProxy("http://www.example.com:8000", verbose = 1)
        c = s.add(1, 2);
        
        Attempts to calculate the "add" function on www.example.com through port
        8000. 
        
        (The expected behavior is that the variable c becomes 3. If there is a
        problem, then an exception may be thrown. See the __init__ function.
    """
    def __init__(self, uri, transport=None, verbose=False, safe_loader = True,
            json_compat = True, throwexcep = True):
        """ This initializes this class as a proxy for a remote server. The
            parameters are:
            
            uri: the connection point on the server using the format.
            
            scheme://host:pathno/target 
            
            (In this implementation, only http and https are supported as 
            schemes.)
            
            transport: An instance of YAMLTransport or YAMLSecureTranspot.
            
            verbose: True only if this instance should print a lot of 
            information to the screen. False otherwise.
            
            safe_loader: the default is True: only allows "simple" YAML data to
            be loaded - sequences, dictionaries and most immutable types. If
            this parameter is False, then arbitrary YAML data can be loaded.
            This can be a security risk. (Only set this parameter to False with
            trusted servers!)
            
            Note: if the transport parameter is not none, then it is modified
            to use the same safe_load settings as this class.
            
            json_compat: if True, all complex objects are sent to the YAML-RPC
            server as dictionaries; no special tagging printed out. If False,
            YAML-RPC requests are delivered with the tag '!yamlrpccall'. The
            default is True.
            
            throwexcep: if True (the default), throw an exception if there is
            a processing error on the remove server.
            
            There is no encoding declaration: it is assumed that all data is
            encoded as UTF-8.
        """

        utype, uri = urllib.splittype(uri)
        if utype not in ("http", "https"):
            raise IOError, "Unsupported YAMLRPC protocol"
        self.__host, self.__handler = urllib.splithost(uri)
        if not self.__handler:
            self.__handler = "/RPC2"
        if transport is None:
            if utype == "https":
                self.__transport = SecureYAMLTransport(safe_loader);
            else:
                self.__transport = YAMLTransport(safe_loader);
        else:
            self.__transport = transport;
            transport.safe_loader = safe_loader;
        self.__verbose = verbose;
        self.json_compat = json_compat;
        self.throwexcep = throwexcep;

    def __request(self, methodname, params):
        """ This calls a method on the remote server. The arguments:
            
            methodname: the name of the method. May contain dots.
            
            params: the parameters as a list or a tuple.
            
            Returns the result if successful, or None if not. May throw
            an exception.
        """
        requestobj = yamlrpcbasic.YAMLRPCCall(_gen_id(), protvers, 
            methodname, params);
        if self.json_compat:
            request = yaml.dump(requestobj.__dict__, default_flow_style=True);
        else:
            request = yaml.dump(requestobj, default_flow_style=True);
        response = self.__transport.request(self.__host, self.__handler, 
            request, self.__verbose);
        errorobj = getAttrWay(response, "error");
        if self.throwexcep and errorobj != None:
            raise YRPCErrorException(errorobj);
        return getAttrWay(response, "result");

    def __repr__(self):
        return ("<YAMLServerProxy for %s%s>" % (self.__host, self.__handler))

    __str__ = __repr__;

    def __getattr__(self, name):
        return _Method(self.__request, name)

    # note: to call a remote object with an non-standard name, use
    # result getattr(server, "strange-python-name")(args)


if __name__ == "__main__":
    s = YAMLServerProxy("http://localhost:8002", None, False, False);
    c = s.add(1, 2);
    print c;
    c = s.system.listMethods();
    print c;
    c = s.system.methodSignature("add");
    print c;
    c = s.system.methodHelp("pow");
    print c;
    c = s.system.methodHelp("add");
    print c;
    
# The following code should cause errors. So we trap them.    
    
    try:
        c = s.bad("other");
        print c;
    except YRPCErrorException, yrpc:
        print yrpc;
    try:
        e = s.echo("foo bar", "baz")
        print e
    except YRPCErrorException, yrpc:
        print yrpc;
    try:
        f = s.echo(5)
        print f
    except YRPCErrorException, yrpc:
        print yrpc;