History

Moved useful methods out into the module body as functions. Added support for more messages, including those that return file information.

David Boddie committed 8 years ago

Commit 4da8c8f1 parent 9f17a3e1 file · +284 −68Browse files

Unified · Split

styx.py+284−68
@@ -45,122 +45,339 @@ class SocketReceiver: · 44 unmodified lines
4545 return data
4646
4747
48−class StyxMessage:
48+def decode_string(stream):
4949
50− def decode_string(self, stream):
51−
52− length = struct.unpack("<H", stream.recv(2))[0]
50+ length = struct.unpack("<H", stream.recv(2))[0]
51+ if length > 0:
5352 return stream.recv(length).decode("utf8")
53+ else:
54+ return u""
5455
56+def encode_string(string):
57+
58+ utf8 = string.encode("utf8")
59+ return struct.pack("<H", len(utf8)) + utf8
60+
61+def decode_data(stream):
62+
63+ n = struct.unpack("<H", stream.recv(2))[0]
64+ return stream.recv(n)
65+
66+def encode_data(data):
67+
68+ return struct.pack("<H", len(data)) + data
69+
70+def decode_qid(stream):
71+
72+ # Store a tuple of values: file/dir, qid version, qid path.
73+ return struct.unpack("<BIQ", stream.recv(13))
74+
75+def encode_qid(qid):
76+
77+ return struct.pack("<BIQ", *qid)
78+
79+
80+class StyxError(Exception):
81+ pass
82+
83+
84+class StyxMessage:
85+
5586 def init(self, size, tag):
5687
5788 self.size = size
5889 self.tag = tag
5990
60− def encode_string(self, string):
61−
62− utf8 = string.encode("utf8")
63− return struct.pack("<H", len(utf8)) + utf8
91+ def decode(self, size, tag, stream):
6492
65− def _parse_format(self, stream):
66−
67− self.fields = {
68− "size": repr(self.size),
69− "tag": repr(self.tag)
70− }
93+ self.init(size, tag)
7194
72− for piece in self.format.split()[3:]:
95+ for name, length in self.format:
7396
74− begin = piece.find("[")
75−
76− if begin != -1:
97+ if length == "s":
98+ # UTF-8 string
99+ value = decode_string(stream)
100+ elif length == "n":
101+ # n bytes of raw data
102+ value = decode_data(stream)
103+ elif length == 2:
104+ value = struct.unpack("<H", stream.recv(2))[0]
105+ elif length == 4:
106+ value = struct.unpack("<I", stream.recv(4))[0]
107+ elif length == 8:
108+ value = struct.unpack("<Q", stream.recv(8))[0]
109+ elif name == "qid":
110+ value = decode_qid(stream)
111+ else:
112+ raise StyxError("Unknown field length specifier: %s" % length)
77113
78− name = piece[:begin]
79− end = piece.find("]", begin)
80− length = piece[begin + 1:end]
81−
82− if length == "s":
83− # UTF-8 string
84− value = self.decode_string(stream)
85− elif length == "n":
86− # n bytes of raw data
87− n = struct.unpack("<H", stream.recv(2))
88− value = stream.recv(n)
89− elif length == "2":
90− value = struct.unpack("<H", stream.recv(2))
91− elif length == "4":
92− value = struct.unpack("<I", stream.recv(4))
93− else:
94− raise
95−
96− self.fields[name] = repr(value)
97− self.__dict__[name] = value
98−
99− def parse(self, size, tag, stream):
100−
101− self.init(size, tag)
102− self._parse_format(stream)
114+ self.__dict__[name] = value
115+
103116 return self
104117
105118 def __repr__(self):
106119
107− return self.repr_format % self.fields
120+ s = self.msg_name + "(tag=" + repr(self.tag)
121+
122+ for name, length in self.format:
123+ s += ", " + "%s=%s" % (name, repr(self.__dict__[name]))
124+
125+ return s + ")"
126+
127+ def encode(self, stream):
128+
129+ # Start with the message type and tag before adding all the other pieces.
130+ data = struct.pack("<bH", self.code, self.tag)
131+
132+ data += self._encode_format(stream)
133+
134+ # Finally, prepend the length to the byte string and send it.
135+ data = struct.pack("<I", len(data) + 4) + data
136+ stream.sendall(data)
137+
138+ def _encode_format(self, stream):
139+
140+ data = ""
141+
142+ for name, length in self.format:
143+
144+ value = self.__dict__[name]
145+
146+ if length == "s":
147+ # UTF-8 string
148+ data += encode_string(value)
149+ elif length == "n":
150+ # n bytes of raw data
151+ data += encode_data(value)
152+ elif length == 2:
153+ data += struct.pack("<H", value)
154+ elif length == 4:
155+ data += struct.pack("<I", value)
156+ elif length == 8:
157+ data += struct.pack("<Q", value)
158+ elif name == "qid":
159+ data += encode_qid(value)
160+
161+ return data
108162
109163
110164 class Tversion(StyxMessage):
111165
166+ msg_name = "Tversion"
112167 code = 100
113− format = "size[4] Tversion tag[2] msize[4] version[s]"
114− repr_format = "Tversion(tag=%(tag)s, msize=%(msize)s, version=%(version)s)"
168+ format = [("msize", 4), ("version", "s")]
115169
116170 def __init__(self, tag = None, msize = 0, version = ""):
117171
118172 self.tag = tag
119173 self.msize = msize
120174 self.version = version
121−
122− def write(self, stream):
123−
124− data = struct.pack("<bHI", self.code, self.tag, self.msize)
125− data += self.encode_string(self.version)
126− data = struct.pack("<I", len(data) + 4) + data
127− stream.sendall(data)
128−
129175
130176 class Rversion(StyxMessage):
131177
178+ msg_name = "Rversion"
132179 code = 101
133− format = "size[4] Rversion tag[2] msize[4] version[s]"
134− repr_format = "Rversion(tag=%(tag)s, msize=%(msize)s, version=%(version)s)"
180+ format = [("msize", 4), ("version", "s")]
135181
136182 def __init__(self, tag = None, msize = 0, version = ""):
137183
138184 self.tag = tag
139185 self.msize = msize
140186 self.version = version
187+
188+
189+class Tattach(StyxMessage):
190+
191+ msg_name = "Tattach"
192+ code = 104
193+ format = [("fid", 4), ("afid", 4), ("uname", "s"), ("aname", "s")]
194+
195+ def __init__(self, tag = None, fid = None, afid = None, uname = None, aname = None):
196+
197+ self.tag = tag
198+ self.fid = fid
199+ self.afid = afid
200+ self.uname = uname
201+ self.aname = aname
202+
203+class Rattach(StyxMessage):
204+
205+ msg_name = "Rattach"
206+ code = 105
207+ format = [("qid", 13)]
208+
209+ def __init__(self, tag = None, qid = None):
210+
211+ self.tag = tag
212+ self.qid = qid
213+
214+
215+class Tstat(StyxMessage):
216+
217+ msg_name = "Tstat"
218+ code = 124
219+ format = [("fid", 4)]
220+
221+ def __init__(self, tag = None, fid = None):
222+
223+ self.tag = tag
224+ self.fid = fid
225+
226+class Rstat(StyxMessage):
227+
228+ msg_name = "Rstat"
229+ code = 125
230+ # Because the format of this message describes the stat information as
231+ # stat[n] it includes a 16-bit length, but the stat object itself
232+ # includes its own length field.
233+ format = [
234+ ("stat_size", 2), ("size", 2), ("type", 2), ("dev", 4), ("qid", 13),
235+ ("mode", 4), ("atime", 4), ("mtime", 4), ("length", 8), ("name", "s"),
236+ ("uid", "s"), ("gid", "s"), ("muid", "s")
237+ ]
141238
142− def write(self, stream):
239+ def __init__(self, tag = None, stat = None):
143240
144− data = struct.pack("<bHI", self.code, self.tag, self.msize)
145− data += self.encode_string(self.version)
241+ self.tag = tag
242+ self.stat = stat
243+
244+ def encode(self, stream):
245+
246+ data = struct.pack("<bH", self.code, self.tag)
247+
248+ # Encode the stat structure.
249+ data += self.stat.encode()
250+
146251 data = struct.pack("<I", len(data) + 4) + data
147252 stream.sendall(data)
148253
149254
255+class Stat:
256+
257+ def __init__(self, type, dev, qid, mode, atime, mtime, length, name, uid,
258+ gid, muid):
259+
260+ self.type = type
261+ self.dev = dev
262+ self.qid = qid
263+ self.mode = mode
264+ self.atime = atime
265+ self.mtime = mtime
266+ self.length = length
267+ self.name = name
268+ self.uid = uid
269+ self.gid = gid
270+ self.muid = muid
271+
272+ def encode(self):
273+
274+ args = (self.type, self.dev) + self.qid + (self.mode, self.atime, self.mtime, self.length)
275+ data = struct.pack("<HIBIQIIIQ", *args)
276+ data += encode_string(self.name)
277+ data += encode_string(self.uid)
278+ data += encode_string(self.gid)
279+ data += encode_string(self.muid)
280+
281+ # Prepend the size of the fields within the stat structure.
282+ data = encode_data(data)
283+
284+ # Prepend the size of the stat structure itself.
285+ return encode_data(data)
286+
287+
288+class Twalk(StyxMessage):
289+
290+ msg_name = "Twalk"
291+ code = 110
292+ format = [("fid", 4), ("newfid", 4), ("nwname", 2)] # nwname of wname[s]
293+
294+ def __init__(self, tag = None, fid = None, newfid = None, wname = None):
295+
296+ self.tag = tag
297+ self.fid = fid
298+ self.newfid = newfid
299+ self.wname = wname
300+
301+ def decode(self, size, tag, stream):
302+
303+ StyxMessage.decode(self, size, tag, stream)
304+
305+ # Decode the names.
306+ self.wname = []
307+ i = 0
308+ while i < self.nwname:
309+ self.wname.append(decode_string(stream))
310+ i += 1
311+
312+ return self
313+
314+ def encode(self, stream):
315+
316+ data = struct.pack("<bH", self.code, self.tag)
317+
318+ self.nwname = len(self.wname)
319+ data += self._encode_format(stream)
320+
321+ # Encode the names separately.
322+ for name in self.wname:
323+ data += encode_string(name)
324+
325+ data = struct.pack("<I", len(data) + 4) + data
326+ stream.sendall(data)
327+
328+
329+class Rwalk(StyxMessage):
330+
331+ msg_name = "Rwalk"
332+ code = 111
333+ format = [("nwqid", 2)] # nwqid of wqid[13]
334+
335+ def __init__(self, tag = None, wqid = None):
336+
337+ self.tag = tag
338+ self.wqid = wqid
339+
340+ def decode(self, size, tag, stream):
341+
342+ StyxMessage.decode(self, size, tag, stream)
343+
344+ # Decode the qids.
345+ self.wqid = []
346+ i = 0
347+ while i < self.nwqid:
348+ self.wqid.append(decode_qid(stream))
349+ i += 1
350+
351+ return self
352+
353+ def encode(self, stream):
354+
355+ data = struct.pack("<bH", self.code, self.tag)
356+
357+ self.nwqid = len(self.wqid)
358+ data += self._encode_format(stream)
359+
360+ # Encode the qids separately.
361+ for qid in self.wqid:
362+ data += encode_qid(qid)
363+
364+ data = struct.pack("<I", len(data) + 4) + data
365+ stream.sendall(data)
366+
150367
151368 MessageTypes = {
152369 Tversion.code: Tversion,
153370 Rversion.code: Rversion,
154371 # Tauth.code: Tauth,
155372 # Rauth.code: Rauth,
156−# Tattach.code: Tattach,
157−# Rattach.code: Rattach,
373+ Tattach.code: Tattach,
374+ Rattach.code: Rattach,
158375 # Terror.code: Terror,
159376 # Rerror.code: Rerror,
160377 # Tflush.code: Tflush,
161378 # Rflush.code: Rflush,
162−# Twalk.code: Twalk,
163−# Rwalk.code: Rwalk,
379+ Twalk.code: Twalk,
380+ Rwalk.code: Rwalk,
164381 # Topen.code: Topen,
165382 # Ropen.code: Ropen,
166383 # Tcreate.code: Tcreate,
@@ -173,14 +390,13 @@ MessageTypes = { · 6 unmodified lines
173390 # Rclunk.code: Rclunk,
174391 # Tremove.code: Tremove,
175392 # Rremove.code: Rremove,
176−# Tstat.code: Tstat,
177−# Rstat.code: Rstat,
393+ Tstat.code: Tstat,
394+ Rstat.code: Rstat,
178395 # Twstat.code: Twstat,
179396 # Rwstat.code: Rwstat
180397 }
181398
182−
183−def parse(sock = None, data = None):
399+def decode(sock = None, data = None):
184400
185401 if sock:
186402 stream = SocketReceiver(sock)
@@ -197,4 +413,4 @@ def parse(sock = None, data = None): · 10 unmodified lines
197413 # Find the relevant message class to handle this type and create an
198414 # instance of it to parse the message data.
199415 Message = MessageTypes[message_type]
200− return Message().parse(size, tag, stream)
416+ return Message().decode(size, tag, stream)