class Connection(object):

def __init__(self, conn):

self._conn = conn

def timeout(self, timeout):

return self._conn.settimeout(timeout)

def close(self):

self._conn.close()

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()

while len(retval) < length:

data = self.recv(length - len(retval))

if len(data) == 0:

return None

retval += data

return retval

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)

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()

self._dir.cleanup()

def next_file(self, initial_timeout=None, timeout=5):

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

fn_len = struct.unpack("B", data)[0]

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]

# Receive the file data

fn = os.path.join(self._dir.name, fn)

with open(fn, "wb+") as f:

while file_len > 0:

data = conn.recv(min(4096, file_len))

f.write(data)

file_len -= len(data)

conn.close()

if file_len is 0:

return fn, mime

else:

os.remove(fn)

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]: