← Back to branch summary

~azzar1/unity/add-show-desktop-key

« back to all changes in this revision

Viewing changes to ivle/chat.py

  • Committer: William Grant
  • Date: 2010-02-24 07:22:43 UTC
  • Revision ID: grantw@unimelb.edu.au-20100224072243-xq5w2we8iuoteen1
Reword and reformat the tour a bit.

Show diffs side-by-side

added added

removed removed

Lines of Context:
ivle/chat.py
21
21
 
22
22
import cjson
23
23
import cStringIO
24
 
import md5
 
24
import hashlib
25
25
import sys
26
26
import os
27
27
import socket
28
28
import traceback
29
29
 
 
30
SOCKETTIMEOUT = 60
 
31
BLOCKSIZE = 1024
 
32
 
30
33
class Terminate(Exception):
31
 
    """Exception thrown when server is to be shut down. It will attempt to sned 
32
 
    the final_response to the client and then exits"""
 
34
    """Exception thrown when server is to be shut down. It will attempt to
 
35
    send the final_response to the client and then exits"""
33
36
    def __init__(self, final_response=None):
34
37
        self.final_response = final_response
35
38
 
36
39
    def __str__(self):
37
40
        return repr(self.final_response)
38
41
 
 
42
class ProtocolError(Exception):
 
43
    """Exception thrown when client violates the the chat protocol"""
 
44
    pass
39
45
 
40
46
def start_server(port, magic, daemon_mode, handler, initializer = None):
41
47
    # Attempt to open the socket.
53
59
            os._exit(0) # kill off parent again.
54
60
        os.umask(077)
55
61
 
 
62
        try:
 
63
            MAXFD = os.sysconf("SC_OPEN_MAX")
 
64
        except:
 
65
            MAXFD = 256
 
66
 
 
67
        # Close all file descriptors, except the socket.
 
68
        for i in xrange(MAXFD):
 
69
            if i == s.fileno():
 
70
                continue
 
71
            try:
 
72
                os.close(i)
 
73
            except OSError:
 
74
                pass
 
75
 
 
76
        si = os.open(os.devnull, os.O_RDONLY)
 
77
        os.dup2(si, sys.stdin.fileno())
 
78
 
 
79
        so = os.open(os.devnull, os.O_WRONLY)
 
80
        os.dup2(so, sys.stdout.fileno())
 
81
 
 
82
        se = os.open(os.devnull, os.O_WRONLY)
 
83
        os.dup2(se, sys.stderr.fileno())
 
84
 
56
85
    if initializer:
57
86
        initializer()
58
87
 
59
88
    while True:
60
89
        (conn, addr) = s.accept()
 
90
        conn.settimeout(SOCKETTIMEOUT)
61
91
        try:
62
92
            # Grab the input
63
 
            buf = cStringIO.StringIO()
64
 
            blk = conn.recv(1024)
65
 
            while blk:
66
 
                buf.write(blk)
67
 
                try:
68
 
                    blk = conn.recv(1024, socket.MSG_DONTWAIT)
69
 
                except:
70
 
                    # Exception thrown if it WOULD block (but we
71
 
                    # told it not to wait) - ie. we are done
72
 
                    blk = None
73
 
            inp = buf.getvalue()
 
93
            inp = recv_netstring(conn)
74
94
            env = cjson.decode(inp)
75
 
            
 
95
 
76
96
            # Check that the message is 
77
 
            digest = md5.new(env['content'] + magic).digest().encode('hex')
 
97
            digest = hashlib.md5(env['content'] + magic).hexdigest()
78
98
            if env['digest'] != digest:
79
99
                conn.close()
80
100
                continue
83
103
 
84
104
            response = handler(content)
85
105
 
86
 
            conn.sendall(cjson.encode(response))
 
106
            send_netstring(conn, cjson.encode(response))
87
107
 
88
108
            conn.close()
89
109
 
90
110
        except Terminate, t:
91
111
            # Try and send final response and then terminate
92
112
            if t.final_response:
93
 
                conn.sendall(cjson.encode(t.final_response))
 
113
                send_netstring(conn, cjson.encode(t.final_response))
94
114
            conn.close()
95
115
            sys.exit(0)
96
116
        except Exception:
103
123
                "value": str(e_val),
104
124
                "traceback": tb_dump.getvalue()
105
125
            }
106
 
            conn.sendall(cjson.encode(json_exc))
 
126
            send_netstring(conn, cjson.encode(json_exc))
107
127
            conn.close()
108
128
 
109
129
 
110
130
def chat(host, port, msg, magic, decode = True):
111
131
    sok = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
112
132
    sok.connect((host, port))
 
133
    sok.settimeout(SOCKETTIMEOUT)
113
134
    content = cjson.encode(msg)
114
 
    digest = md5.new(content + magic).digest().encode("hex")
 
135
    digest = hashlib.md5(content + magic).hexdigest()
115
136
    env = {'digest':digest,'content':content}
116
 
    sok.send(cjson.encode(env))
117
 
 
118
 
    buf = cStringIO.StringIO()
119
 
    blk = sok.recv(1024)
120
 
    while blk:
121
 
        buf.write(blk)
122
 
        try:
123
 
            blk = sok.recv(1024, socket.MSG_DONTWAIT)
124
 
        except socket.error, e:
125
 
            if e[0] != 11:
126
 
                raise
127
 
            # Exception thrown if it WOULD block (but we
128
 
            # told it not to wait) - ie. we are done
129
 
            blk = None
130
 
    inp = buf.getvalue()
 
137
    json = cjson.encode(env)
 
138
 
 
139
    send_netstring(sok, json)
 
140
    inp = recv_netstring(sok)
 
141
 
131
142
    sok.close()
132
143
 
133
144
    if decode:
135
146
    else:
136
147
        return inp
137
148
 
 
149
 
 
150
def send_netstring(sok, data):
 
151
    netstring = "%d:%s,"%(len(data),data)
 
152
    sok.sendall(netstring)
 
153
 
 
154
 
 
155
def recv_netstring(sok):
 
156
    # Decode netstring
 
157
    size_buffer = []
 
158
    c = sok.recv(1)
 
159
    while c != ':':
 
160
        # Limit the Netstring to less than 10^10 bytes (~1GB).
 
161
        if len(size_buffer) >= 10:
 
162
            raise ProtocolError(
 
163
                    "Could not read Netstring size in first 9 bytes: '%s'"%(
 
164
                    ''.join(size_buffer)))
 
165
        size_buffer.append(c)
 
166
        c = sok.recv(1)
 
167
    try:
 
168
        # Message should be length plus ',' terminator
 
169
        recv_expected = int(''.join(size_buffer)) + 1
 
170
    except ValueError, e:
 
171
        raise ProtocolError("Could not decode Netstring size as int: '%s'"%(
 
172
                ''.join(size_buffer)))
 
173
 
 
174
    # Read data
 
175
    buf = []
 
176
    recv_data = sok.recv(min(recv_expected, BLOCKSIZE))
 
177
    recv_size = len(recv_data)
 
178
    while recv_size < recv_expected:
 
179
        buf.append(recv_data)
 
180
        recv_data = sok.recv(min(recv_expected-recv_size, 1024))
 
181
        recv_size = recv_size + len(recv_data)
 
182
    assert(recv_size == recv_expected)
 
183
 
 
184
    # Did we receive the correct amount?
 
185
    if recv_data[-1] != ',':
 
186
        raise ProtocolError("Netstring did not end with ','")
 
187
    buf.append(recv_data[:-1])
 
188
 
 
189
    return ''.join(buf)