diff --git a/mailcap-client.py b/mailcap-client.py index 970762e..6405ef3 100644 --- a/mailcap-client.py +++ b/mailcap-client.py @@ -5,6 +5,7 @@ import os import sys import mimetypes import subprocess +import hashlib s = socket.socket(socket.AF_UNIX) s.connect(os.path.join(os.environ["HOME"], ".mailcap.sock")) @@ -19,4 +20,11 @@ data = open(filename, "rb").read() def strtr(x, c="B"): return struct.pack(c, len(x)) + x -s.sendall(strtr(filename.split(os.pathsep)[-1].encode('utf-8')) + strtr(mime.encode('utf-8')) + strtr(data, "!I")) +hash = hashlib.sha256() +hash.update(data) + +s.sendall(strtr(filename.split(os.pathsep)[-1].encode('utf-8')) + strtr(mime.encode('utf-8')) + hash.digest()) + +response = s.recv(1) +if len(response) and response[0] == 1: + s.sendall(strtr(data, "!I")) diff --git a/mailcap-server.py b/mailcap-server.py index db16652..c3f2603 100644 --- a/mailcap-server.py +++ b/mailcap-server.py @@ -7,6 +7,7 @@ import mailcap import sys import argparse import daemon +import hashlib sockpath = "mailcap.sock" @@ -27,6 +28,9 @@ class Server(object): def recv(self, size): data = self._conn.recv(size) return data + + def send(self, data): + return self._conn.send(data) def recv_exact(self, length): retval = bytes() @@ -37,11 +41,16 @@ class Server(object): retval += data return retval - def __init__(self, path="mailcap.sock"): + def __init__(self, path="mailcap.sock", dir = None): self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) self._socket.bind(path) self._socket.listen(3) - self._dir = tempfile.TemporaryDirectory() + if dir: + if not os.path.exists(dir): + os.makedirs(dir) + self._dir = dir + else: + self._dir = tempfile.TemporaryDirectory() def close(self): self._socket.close() @@ -51,6 +60,7 @@ class Server(object): conn = self.Connection(self._socket.accept()[0]) conn.timeout(initial_timeout) + # Receive filename length data = conn.recv_exact(1) if data is None: return None @@ -58,31 +68,64 @@ class Server(object): conn.timeout(timeout) + # Receive filename data = conn.recv_exact(fn_len) if data is None: return None fn = data.decode('utf-8') + + # Strip the filename of /, do sanity checks fn = fn.replace("/", "_") if len(fn) == 0 or fn in ['.', '..']: + s.close() return None - + + # Receive the mime type length data = conn.recv_exact(1) if data is None: return None mime_len = struct.unpack("B", data)[0] + # Receive the mime type data = conn.recv_exact(mime_len) if data is None: return None mime = data.decode('utf-8') + # Receive the SHA256 hash + data = conn.recv_exact(hashlib.sha256().digest_size) + if data is None: + return None + + full_name = os.path.join(self._dir.name, fn) + append = 0 + # While the filename already exists, check if we have a hash match + while os.path.exists(full_name): + hash = hashlib.sha256() + hash.update(open(full_name, "rb").read()) + + # Found a matching file. Do not receive the file. + if data == hash.digest(): + conn.send(b'\x00') + conn.close() + return full_name, mime + + append += 1 + hash = hashlib.sha256() + name_parts = fn.split('.') + if len(name_parts) > 1: + full_name = os.path.join(self._dir.name, "%s-%d.%s" % ('.'.join(name_parts[:-1]), append, name_parts[-1])) + else: + full_name = os.path.join(self._dir.name, "%s-%d" % (fn, append, name_parts[-1])) + conn.send(b'\x01') + + # Receive the file length data = conn.recv_exact(4) if data is None: return None file_len = struct.unpack("!I", data)[0] - print(fn, mime) - + # Receive the file data fn = os.path.join(self._dir.name, fn) with open(fn, "wb+") as f: while file_len > 0: @@ -98,12 +141,26 @@ class Server(object): os.remove(fn) return None +def run_server(args): + sv = Server(path=args.socket, dir=args.dir) + caps = mailcap.getcaps() + open(args.pid, "w+").write(str(os.getpid())) + while True: + rv = sv.next_file() + if rv: + f, mime = rv + print(f, mime) + match = mailcap.findmatch(caps, mime.split(';')[0], filename=f) + if match and match[0]: + os.system(match[0]) + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--pid', '-p', default=os.path.join(os.environ['HOME'], '.mailcap.pid')) parser.add_argument('--action', '-a', default='run', choices=['run', 'kill']) parser.add_argument('--no-daemonize', '-n', action='store_true') parser.add_argument('--socket', '-s', default=os.path.join(os.environ['HOME'], '.mailcap.sock')) + parser.add_argument('--dir', '-d', default=None) args = parser.parse_args() if args.action == 'kill': @@ -125,17 +182,8 @@ if __name__ == "__main__": if os.path.exists(args.socket): os.remove(args.socket) - with daemon.DaemonContext(detach_process=not args.no_daemonize, files_preserve=[1,2]): - sv = Server(args.socket) - caps = mailcap.getcaps() - open(args.pid, "w+").write(str(os.getpid())) - while True: - rv = sv.next_file() - if rv: - f, mime = rv - print(f, mime) - match = mailcap.findmatch(caps, mime.split(';')[0], filename=f) - if match and match[0]: - os.system(match[0]) - if 'nodelete' not in match[1]: - os.remove(f) + if args.no_daemonize: + run_server(args) + else: + with daemon.DaemonContext(): + run_server()