Ticket #45: SimpleYAMLRPCServer.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
# SimpleYAMLRPCServer.py
"""Simple YAML-RPC Server.

This module can be used to create simple YAML-RPC servers
by creating a server and either installing functions, a
class instance, or by extending the SimpleYAMLRPCServer
class.

It can also be used to handle YAML-RPC requests in a CGI
environment using CGIYAMLRPCRequestHandler.

A list of possible usage patterns follows:

1. Install functions:

server = SimpleYAMLRPCServer(("localhost", 8000))
server.register_function(pow)
server.register_function(lambda x,y: x+y, 'add')
server.serve_forever()

2. Install an instance:

class MyFuncs:
    def __init__(self):
        # make all of the string functions available through
        # string.func_name
        import string
        self.string = string
    def _listMethods(self):
        # implement this method so that system.listMethods
        # knows to advertise the strings methods
        return list_public_methods(self) + \
                ['string.' + method for method in list_public_methods(self.string)]
    def pow(self, x, y): return pow(x, y)
    def add(self, x, y) : return x + y

server = SimpleYAMLRPCServer(("localhost", 8000))
server.register_introspection_functions()
server.register_instance(MyFuncs())
server.serve_forever()

3. Install an instance with custom dispatch method:

class Math:
    def _listMethods(self):
        # this method must be present for system.listMethods
        # to work
        return ['add', 'pow']
    def _methodHelp(self, method):
        # this method must be present for system.methodHelp
        # to work
        if method == 'add':
            return "add(2,3) => 5"
        elif method == 'pow':
            return "pow(x, y[, z]) => number"
        else:
            # By convention, return empty
            # string if no help is available
            return ""
    def _dispatch(self, method, params):
        if method == 'pow':
            return pow(*params)
        elif method == 'add':
            return params[0] + params[1]
        else:
            raise 'bad method'

server = SimpleXMLRPCServer(("localhost", 8000))
server.register_introspection_functions()
server.register_instance(Math())
server.serve_forever()

4. Subclass SimpleYAMLRPCServer:

class MathServer(SimpleYAMLRPCServer):
    def _dispatch(self, method, params):
        try:
            # We are forcing the 'export_' prefix on methods that are
            # callable through YAML-RPC to prevent potential security
            # problems
            func = getattr(self, 'export_' + method)
        except AttributeError:
            raise Exception('method "%s" is not supported' % method)
        else:
            return func(*params)

    def export_add(self, x, y):
        return x + y

server = MathServer(("localhost", 8000))
server.serve_forever()

5. CGI script:

server = CGIYAMLRPCRequestHandler()
server.register_function(pow)
server.handle_request()
"""

# This implementation is based on SimpleJSONRPCServer, which was
# was converted from SimpleXMLRPCServer by David McNab 
# (david@rebirthing.co.nz). Current implementation written by Peter
# Murphy.

# Original SimpleXMLRPCServer module was written by Brian
# Quinlan (brian@sweetapp.com), Based on code written by Fredrik Lundh.

import xmlrpclib
from xmlrpclib import Fault
import SocketServer
import BaseHTTPServer
import sys
import os

import SimpleXMLRPCServer
import yaml
from yamlrpcbasic import *

import traceback

class SimpleYAMLRPCDispatcher(SimpleXMLRPCServer.SimpleXMLRPCDispatcher):
    """ Mix-in class that dispatches YAML-RPC requests. Based on 
        SimpleXMLRPCDispatcher, but overrides marshaled_dispatch for 
        YAML-RPC.

        This class is used to register YAML-RPC method handlers and then to
        dispatch them. There should never be any reason to instantiate this
        class directly.
    """

    def __init__(self, safe_loader = True, json_compat = True):
        """ The parameters are:
            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 clients!)
            
            json_compat: if True, all complex objects are sent to YAML-RPC
            clients as dictionaries; no special tagging printed out. If False,
            YAML-RPC responses are delivered with the tag '!yamlrpcret'. The
            default is True.
        """
        self.json_compat = json_compat;
        self.loadfunc(safe_loader);
        SimpleXMLRPCServer.SimpleXMLRPCDispatcher.__init__(self, True, "utf-8");

    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 _marshaled_dispatch(self, data, dispatch_method = None):
        """ Dispatches a YAML-RPC method from marshalled (YAML) data.
    
            YAML-RPC methods are dispatched from the marshalled (YAML) data
            using the _dispatch method and the result is returned as
            marshalled data. For backwards compatibility, a dispatch
            function can be provided as an argument (see comment in
            SimpleYAMLRPCRequestHandler.do_POST) but overriding the
            existing method through subclassing is the prefered means of 
            changing method dispatch behavior.
        
            The "data" part can be one of the following:
            (a) An object with tag '!yamlrpccall'.
            (b) A dictionary/YAML map (with the same keys). This is to allow 
            backward compatibilty with JSON-RPC.
        
        """
        rawreq = self._loadfunc(data)
        id = getAttrWay(rawreq, "id", None);
        method = getAttrWay(rawreq, "method", None);
        params = getAttrWay(rawreq, "params", []);
        OurResponseObj = YAMLRPCReturn(id, "1.0", None, None);
        
    
        # We generate a response.
        try:
            if dispatch_method is not None:
                response = dispatch_method(method, params)
            else:
                response = self._dispatch(method, params)
            OurResponseObj.result = response
#        except Fault, fault:
            #response = xmlrpclib.dumps(fault)
 #           OurResponseObj.error = YAMLRPCError("YAMLRPCError", 666, repr(response), None);
        except:
            OurResponseObj.error = YAMLRPCError("YAMLRPCError", 667, 
                "%s:%s" % (sys.exc_type, sys.exc_value), None);

# The following code is for JSON compatibility.

        if self.json_compat:
            OurResponseObj = OurResponseObj.__dict__;
            OurError = OurResponseObj["error"];
            if OurError != None:
                OurResponseObj["error"] = OurError.__dict__;
        return yaml.dump(OurResponseObj);
    

class SimpleYAMLRPCRequestHandler(SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
    """Simple YAML-RPC request handler class.

    Handles all HTTP POST requests and attempts to decode them as
    XML-RPC requests.
    """
    def do_POST(self):
        """Handles the HTTP POST request.
    
        Attempts to interpret all HTTP POST requests as YAML-RPC calls,
        which are forwarded to the server's _dispatch method for handling.
        """
        try:
            # get arguments
            data = self.rfile.read(int(self.headers["content-length"]))
            # In previous versions of SimpleXMLRPCServer, _dispatch
            # could be overridden in this class, instead of in
            # SimpleXMLRPCDispatcher. To maintain backwards compatibility,
            # check to see if a subclass implements _dispatch and dispatch
            # using that method if present.
            response = self.server._marshaled_dispatch(
                    data, getattr(self, '_dispatch', None)
                )
        except: # This should only happen if the module is buggy
            # internal error, report as HTTP server error
            self.send_response(500)
            self.end_headers()
            raise;
        else:
            # got a valid XML RPC response
            self.send_response(200)
            self.send_header("Content-type", "text/yaml")
            self.send_header("Content-length", str(len(response)))
            self.end_headers()
            self.wfile.write(response)
    
            # shut down the connection
            self.wfile.flush()
            self.connection.shutdown(1)
    
class SimpleYAMLRPCServer(SocketServer.TCPServer,
                         SimpleYAMLRPCDispatcher):
    """Simple YAML-RPC server.

    Simple YAML-RPC server that allows functions and a single instance
    to be installed to handle requests. The default implementation
    attempts to dispatch YAML-RPC calls to the functions or instance
    installed in the server. Override the _dispatch method inhereted
    from SimpleYAMLRPCDispatcher to change this behavior.
    """
    def __init__(self, addr, requestHandler=SimpleYAMLRPCRequestHandler,
                 logRequests=1, safe_loader = True, json_compat = True):
        self.logRequests = logRequests

        SimpleYAMLRPCDispatcher.__init__(self, safe_loader, json_compat)
        SocketServer.TCPServer.__init__(self, addr, requestHandler)

class CGIYAMLRPCRequestHandler(SimpleYAMLRPCDispatcher):
    """Simple handler for YAML-RPC data passed through CGI."""
    
    def __init__(self, safe_loader = True, json_compat = True):
        SimpleYAMLRPCDispatcher.__init__(safe_loader, json_compat)
    
    def handle_get(self):
        """Handle a single HTTP GET request.
    
        Default implementation indicates an error because
        XML-RPC uses the POST method.
        """
    
        code = 400
        message, explain = \
                 BaseHTTPServer.BaseHTTPRequestHandler.responses[code]
    
        response = BaseHTTPServer.DEFAULT_ERROR_MESSAGE % \
            {
             'code' : code,
             'message' : message,
             'explain' : explain
            }
        print 'Status: %d %s' % (code, message)
        print 'Content-Type: text/html'
        print 'Content-Length: %d' % len(response)
        print
        sys.stdout.write(response)
    
    def handle_request(self, request_text = None):
        """Handle a single YAML-RPC request passed through a CGI post method.
    
        If no YAML data is given then it is read from stdin. The resulting
        YAML-RPC response is printed to stdout along with the correct HTTP
        headers.
        """
        if request_text is None and \
            os.environ.get('REQUEST_METHOD', None) == 'GET':
            self.handle_get()
        else:
            # POST data is normally available through stdin
            if request_text is None:
                request_text = sys.stdin.read()
    
            self.handle_yamlrpc(request_text)
    
    def handle_yamlrpc(self, request_text):
        """Handle a single YAML-RPC request"""
    
        response = self._marshaled_dispatch(request_text)
    
        print 'Content-Type: text/yaml'
        print 'Content-Length: %d' % len(response)
        print
        sys.stdout.write(response)
    
if __name__ == '__main__':
    server = SimpleYAMLRPCServer(("localhost", 8002), safe_loader = False, json_compat = True)
    server.register_function(pow)
    server.register_function(lambda x,y: x+y, 'add');
    server.register_function(lambda x: x, u"echo");
    server.register_introspection_functions();
    server.serve_forever()