websred.py

!/usr/bin/python
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-d', '--daemonize', dest='daemon', action='store_true')
parser.add_argument('-p', '--port',      dest='port', type=int, default=8079)
parser.add_argument('-i', '--subscribe', dest='subscribe')
parser.add_argument('-o', '--publish',   dest='publish')
parser.add_argument('-v', '--debug',     dest='debug', action='store_true')
parser.add_argument('-l', '--log',       dest='log', action='store_true')

args = parser.parse_args()

if args.log:
    import sys
    from twisted.python import log
    log.startLogging(sys.stdout)

from twisted.internet import reactor, protocol, defer
from txws import WebSocketFactory
import txredisapi as redis


class MyWebSocketProtocol(protocol.Protocol):

    def connectionLost(self,reason=None):
        del self.factory.clients[self.transport.getPeer()]

    def dataReceived(self, msg):
        if args.debug: print('WEB', self.transport.getPeer(), msg)
        bridge.webtoredis(msg)

class MyWebSocketFactory(protocol.Factory):
    clients = {}

    def __init__(self):
        bridge.web = self

    def buildProtocol(self, addr):
      if args.debug: print('%s connected to web socket' % addr)
      a = MyWebSocketProtocol()
      a.factory = self
      self.clients[addr] = a
      return a

    def sendToAll(self, msg):
      for a in self.clients.values():
          a.transport.write(msg)

class MyRedisProtocol(redis.SubscriberProtocol):
    def connectionMade(self):
        if args.debug: print('Connected to Redis subscriber')
        self.subscribe(args.subscribe)

    def messageReceived(self, channel, pattern, message):
        message = str(message)
        if args.debug: print('RED', message)
        bridge.redistoweb(message)

    def connectionLost(self, reason):
        pass #print "lost connection:", reason

class MyRedisFactory(redis.SubscriberFactory):

    maxDelay = 120
    continueTrying = True
    protocol = MyRedisProtocol

class Bridge:
    web = None
    red = None

    def redistoweb(self, message):
        if not self.web: return
        self.web.sendToAll(message)

    def webtoredis(self, message):
        if not self.red: return
        self.red.publish(args.publish, message)

def fork(func):
    import os
    pid = os.fork()
    if pid > 0: return
    os.chdir("/")
    os.setsid()
    os.umask(0)
    pid = os.fork()
    if pid > 0: return
    func()
    os._exit(os.EX_OK)

@defer.inlineCallbacks
def redconnect():
    red = yield redis.Connection('localhost', 6379)
    if args.debug: print('Connected to Redis')
    bridge.red = red

bridge = Bridge()

def main():
    reactor.listenTCP(args.port, WebSocketFactory(MyWebSocketFactory()))
    if args.subscribe:
        reactor.connectTCP("localhost", 6379, MyRedisFactory())
    if args.publish:
        reactor.callWhenRunning(redconnect)
    reactor.run()

if args.daemon:
    fork(main)
else:
    main()