websocket.py

gistfile1.txt

#!/usr/bin/env python
#-*- coding:utf-8 -*-
from optparse import OptionParser
import SocketServer
import struct
import hashlib
import itertools

class WebsocketHandler(SocketServer.BaseRequestHandler):
    def handle(self):
        '''クライアントとTCP接続したときに呼び出される。'''
        self.CLIENT_TERMINATED =False

        # Handshakeを確立する
        self.handshake(self.request.recv(1024))

        # クライアントとの相互通信
        while True:
            data = self.request.recv(1024)
            if len(data):
                self.checkDataFrame(data)
            else:
                # Chrome 6用の処理
                self.CLIENT_TERMINATED = True

            if self.CLIENT_TERMINATED:
                print '/client terminated/フラグがセットされたのでコネクションを切断します'
                self.closeConnection()
                break
            else:
                print '[%s:%s] [%s]' % (self.client_address[0], self.client_address[1], self.RAW_DATA)
                if len(self.RAW_DATA) != 0:
                    self.request.send('\x00' + self.RAW_DATA + '\xFF')
                    print self.RAW_DATA

    def checkDataFrame(self, data):
        '''ハンドシェイク後に受信したデータを解釈する。

        受信したデータフレームの/type/を調べtext-frameかclosing-frameを
        判断してデータをRAW_DATAに入れる。
        '''
        if data[0] != '\x00':
            self.CLIENT_TERMINATED = True
        else:  # text-frameの場合
            xs = itertools.takewhile(lambda x : ord(x) != 0xFF, data[1:])
            self.RAW_DATA = "".join(list(xs))

    def closeConnection(self):
        '''クライアントに対してclosing-frameを送信する。'''
        print '%s:%s にclosing-frameを送信' % (self.client_address[0], self.client_address[1])
        self.request.send('\xFF\x00')
        self.finish()

    def handshake(self, data):
        '''クライアントとハンドシェイクを確立する。'''
        print '<' * 10, '%s:%s からハンドシェイクデータを受信した' % (self.client_address[0], self.client_address[1])
        print data
        print '<' * 40

        (fields, key) = self.parse(data)
        print (fields, key)
        self.sendHandshake(fields, key)

    def parse(self, data):
        '''クライアントから受信したハンドシェイクデータを解析する'''
        # 2行目以降の連続したフィールドをバラす
        fields = data.split('\r\n')
        blankLine = False
        d = dict()
        key = ''
        for field in fields[1:]:
            if field == '':
                blankLine = True
                continue
            if not blankLine:
                fieldName, fieldValue = field.split(': ', 1)
                d[fieldName.lower()] = fieldValue  # フィールド名は大文字と小文字を区別しない
            else:
                key = field
                break
        return (d, key)

    def decode(self, s):
        '''ハンドシェイク中のキーを解析する。'''
        n = filter(lambda c : c.isdigit(), s)
        m = filter(lambda c : c == ' '   , s)
        return int(n) / len(m)

    def sendHandshake(self, fields, key):
        '''クライアントにハンドシェイクを送信する。'''
        # Sec-WebSocket-Locationフィールドの値を作成
        LOCATION = 'wss://' if SECURE_FLAG else 'ws://'
        LOCATION += HOST
        LOCATION += ':' + str(PORT) if PORT else ''
        LOCATION += RESOURCE_NAME

        # チャレンジレスポンス
        part1 = self.decode(fields['sec-websocket-key1'])
        part2 = self.decode(fields['sec-websocket-key2'])
        CHALLENGE = struct.pack('>I', part1)  # 値を32bitのビッグエンディアンのバイナリーにする
        CHALLENGE += struct.pack('>I', part2) # 値を32bitのビッグエンディアンのバイナリーにする
        CHALLENGE += key
        RESPONSE = hashlib.md5(CHALLENGE).digest() # /chalenge/のMD5 fingerprintを/response/に入れる

        # 送信するハンドシェイクデータの作成
        handshakeData = 'HTTP/1.1 101 WebSocket Protocol Handshake\r\n'
        handshakeData += 'Upgrade: WebSocket\r\n'
        handshakeData += 'Connection: Upgrade\r\n'
        handshakeData += 'Sec-WebSocket-Origin: ' + ORIGIN + '\r\n'
        handshakeData += 'Sec-WebSocket-Location: ' + LOCATION + '\r\n'
        if PROTOCOL:
            handshakeData += 'Sec-WebSocket-Protocol: ' + PROTOCOL + '\r\n'
        handshakeData += '\r\n' + RESPONSE
        self.request.send(handshakeData)  # ハンドシェイクデータの送信

if __name__ == "__main__":
    usage = u'%prog [-p ポート番号] オリジン [-s サブプロトコル名] [-r リソース]'
    parser = OptionParser(usage=usage)
    parser.add_option('-p', '--port', dest='port', type='int', default=8080,help=u'ポート番号(デフォルトは8080)')
    parser.add_option('-s', '--subprotocol', dest='subprotocol', help=u'サブプロトコル名')
    parser.add_option('-r', '--resource', dest='resource', default='/', help=u'リソース')
    options, args = parser.parse_args()
    if len(args) < 1:
        parser.error('引数を1つ入力してください')
    elif len(args) > 1:
        parser.error('引数が多いです')

    HOST = 'localhost'
    PORT = options.port
    ORIGIN = args[0]
    PROTOCOL = options.subprotocol
    RESOURCE_NAME = options.resource
    SECURE_FLAG = False

    SocketServer.ThreadingTCPServer.allow_reuse_address = True
    server = SocketServer.ThreadingTCPServer((HOST, PORT), WebsocketHandler)

    print 'Ctrl-cで終了します'
    server.serve_forever()