knock.py

#: vim set encoding=utf-8 :
##
 # Knock
 # Simple RPC for Python
 #
 # author 0x77
 # version 0.3
##

# Imports
import sys
import errno

import hmac
import hashlib

from gevent import socket
from gevent.server import StreamServer

try:
    import simplejson
except ImportError:
    import json

import jsonpickle

# Settings
LOGGING     = False
SIGN_SEP    = chr(127)
SIGN_LENGTH = 41
MSG_SIZE    = 8
MSG_END     = chr(0)

# Globals
class Global(object):
    def __init__(self, items):
        for item in items:
            setattr(self, item, None)

g = Global(['secret', 'client', 'logging', 'beforeJSON'])

# Private
_is_client = False
_is_server = False
_methods   = {}
_cmethods  = []

# Messages
MSG_OK   = b'KNOCKGUD'
MSG_FAIL = b'KNOCKBAD'
MSG_HI   = b'KNOCKHAI'
MSG_CALL = b'KNOCKRPC'
MSG_BYE  = b'KNOCKBYE'

# Utils
class RPCCaller(object):
    def __init__(self, methods):
        self.methods = methods

    def __getattr__(self, name):
        # Redirect function calls to RPC
        if name in self.methods:
            def caller(*args, **kwargs):
                return self._rpc_call(name, *args, **kwargs)

            return caller

    def close(self):
        msg_send = _msg_send

        msg_send(g.client, MSG_BYE)
        g.client.shutdown(2)
        g.client.close()

        self = None

    def _rpc_call(self, method, *args, **kwargs):
        msg_send = _msg_send
        msg_recv = _msg_recv

        # Send an RPC request
        msg_send(g.client, MSG_CALL, True)

        # Send the method name
        msg_send(g.client, method)

        # Check for valid method
        response = msg_recv(g.client, MSG_SIZE)
        if response != MSG_OK:
            msg_send(g.client, MSG_BYE, True)
            raise EOFError('Knock: Error fetching the results from "%s".' % (method))

        # Send *args and **kwargs
        msg_send(g.client, args, obj=True)
        msg_send(g.client, kwargs, obj=True)

        # Parse result
        result_type = msg_recv(g.client)

        if result_type in ['list', 'dict', 'instance']:
            result = msg_recv(g.client, obj=True, decode=False)
        else:
            result = msg_recv(g.client)

        # print(result_type)

        if result_type == 'int':
            result = int(result)
        elif result_type == 'float':
            result = float(result)
        elif result_type == 'str':
            result = str(result)
        elif result_type == 'instance':
            if type(result).__name__ == 'dict':
                raise StandardError('Knock: Trying to instance a non-existent class.')
        elif result_type not in ['list', 'dict']:
            raise TypeError('Knock: Argument type is not supported: "%s".' % (result_type,))

        return result

def _msg_recv(conn, size=None, obj=False, decode=True):
    message = b''

    if size == None:
        while True:
            c = conn.recv(1)
            if c == MSG_END:
                break

            message += c
    else:
        message = conn.recv(size + SIGN_LENGTH)

    if len(message) == 0:
        return MSG_BYE

    if decode:
        message = message.decode()

    r_digest, data = message.split(SIGN_SEP)
    digest = hmac.new(g.secret, data, hashlib.sha1).hexdigest()

    if r_digest != digest:
        raise StandardError('Knock: Integrity check failed.')
    else:
        if obj:
            # return pickle.loads(data, -1)
            # return json.loads(data)
            return jsonpickle.decode(data)
        else:
            return data

def _msg_send(conn, message, fixed=False, obj=False):
    if obj:
        if g.beforeJSON != None:
            message = g.beforeJSON(message)

        data = jsonpickle.encode(message)
    else:
        data = str(message).encode()

    if not fixed:
        end = MSG_END
    else:
        if g.logging:
            print('< '+ message)

        end = b''

    digest = hmac.new(g.secret, data, hashlib.sha1).hexdigest()
    digest = str(digest).encode()

    conn.send(digest + SIGN_SEP + data + end)

# Methods
def handle_echo(server, address):
    LOGGING = g.logging

    msg_recv = _msg_recv
    msg_send = _msg_send

    # conn, _ = server.accept()
    conn = server

    try:
        while True:
            msg = msg_recv(conn, MSG_SIZE)

            if LOGGING:
                print('> '+ msg)

            if msg == MSG_HI:
                chk_secret = msg_recv(conn)

                if chk_secret == g.secret:
                    msg_send(conn, MSG_OK, True)
                    msg_send(conn, _cmethods, obj=True)
                else:
                    msg_send(conn, MSG_FAIL, True)
                    conn.close()
                    break
            elif msg == MSG_CALL:
                method = msg_recv(conn)

                if LOGGING:
                    print('-> '+ method)

                if method in _methods:
                    msg_send(conn, MSG_OK, True)

                    args = msg_recv(conn, obj=True, decode=False)
                    kwargs = msg_recv(conn, obj=True, decode=False)

                    if LOGGING:
                        print('-> args = '+ str(args))
                        print('-> kwargs = '+ str(kwargs))

                    result = _methods[method]['method'](*args, **kwargs)
                    result_type = type(result).__name__

                    if result_type == 'list' or result_type == 'dict':
                        msg_send(conn, result_type)
                        msg_send(conn, result, obj=True)
                    elif result_type not in ['int', 'str', 'float']:
                        msg_send(conn, 'instance')
                        msg_send(conn, result, obj=True)
                    else:
                        msg_send(conn, result_type)
                        msg_send(conn, result)

                    if LOGGING:
                        print('<- '+ str(result))
                else:
                    msg_send(conn, MSG_FAIL, True)
                    conn.close()
                    break

            elif msg == MSG_BYE:
                conn.close()
                break
    except socket.error as e:
        if e.errno == errno.EPIPE:
            if LOGGING:
                print('BROKEN_PIPE')

            conn.close()

def serve(secret, logging=False, before_json=None):
    if _is_client:
        raise StandardError("Knock is already connected.")

    g.secret = secret
    g.logging = logging
    g.beforeJSON = before_json

    # try:
    # Initialize the server
    server = StreamServer(('0.0.0.0', 1427), handle_echo)
    server.serve_forever()
    # except:
    #     print('Knock: Socket error.')
    #     sys.exit(1)

def connect(ip, secret):
    if _is_server:
        raise StandardError("Knock is already connected.")

    msg_recv = _msg_recv
    msg_send = _msg_send

    g.secret = secret

    # try:
    g.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    g.client.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    g.client.connect((ip, 1427))

    msg_send(g.client, MSG_HI, True)
    msg_send(g.client, secret)

    response = msg_recv(g.client, MSG_SIZE)
    if response == MSG_OK:
        methods = msg_recv(g.client, obj=True, decode=False)
        return RPCCaller(methods)
    else:
        print('Knock: Handshake failed.')
        sys.exit(1)
    # except:
    #     print('Knock: Socket error.')
    #     sys.exit(1)

def expose(alias=None):
    def decorator(func):
        def wrapper():
            if sys.version_info > (2,5):
                f_alias = alias or func.__name__
            else:
                f_alias = alias or func.func_name

            function = func

            _methods[f_alias] = { 'method': function }
            _cmethods.append(f_alias)
        return wrapper()
    return decorator