Add a public mode.

In this mode, the server returns short easy to remember URLs.

1 files changed, 194 insertions, 14 deletions

diff --git a/fshare.py b/fshare.py
index 37c36bd..5e493ab 100755
--- a/fshare.py
+++ b/fshare.py
@@ -22,6 +22,7 @@ import argparse
 import contextlib
 import fcntl
 import hmac
+import itertools
 import json
 import os
 import os.path
@@ -34,6 +35,7 @@ import shutil
 import socket
 import sys
 import tempfile
+import threading
 from urllib import parse as urlparse
 
 # TODO: detect and store mime types
@@ -44,19 +46,17 @@ def _open_dirfd(path, mode = 'r', perms = 0o644, dir_fd = None, **kwargs):
     and operation with respect to a directory FD.
     """
     flags = 0
-
     if '+' in mode:
-        flags |= os.O_RDWR
+        flags = os.O_RDWR
+    elif 'r' in mode:
+        flags = os.O_RDONLY
     else:
-        if 'r' in mode:
-            flags |= os.O_RDONLY
-        elif 'w' in mode or 'a' in mode:
-            flags |= os.O_WRONLY
+        flags = os.O_WRONLY
 
     if 'w' in mode:
         flags |= os.O_CREAT | os.O_TRUNC
     elif 'a' in mode:
-        flags |= os.O_CREAT | os.O_TRUNC
+        flags |= os.O_CREAT | os.O_APPEND
 
     fdopen_mode = mode
     if 'x' in mode:
@@ -70,17 +70,165 @@ def _open_dirfd(path, mode = 'r', perms = 0o644, dir_fd = None, **kwargs):
         os.close(fd)
         raise
 
+class UrlEncoder(object):
+    """
+    Author: Michael Fogleman
+    License: MIT
+    Link: http://code.activestate.com/recipes/576918/
+    """
+
+    _alphabet   = 'mn6j2c4rv8bpygw95z7hsdaetxuk3fq'
+    _block_size = 8
+    _min_length = 1
+
+    def __init__(self, block_size = None):
+        self._mask       = (1 << self._block_size) - 1
+        self._mapping    = range(self._block_size)
+
+    def encode_url(self, n):
+        return self.enbase(self.encode(n))
+
+    def decode_url(self, n):
+        return self.decode(self.debase(n))
+
+    def encode(self, n):
+        return (n & ~self._mask) | self._encode(n & self._mask)
+
+    def _encode(self, n):
+        result = 0
+        for i, b in enumerate(reversed(self._mapping)):
+            if n & (1 << i):
+                result |= (1 << b)
+        return result
+
+    def decode(self, n):
+        return (n & ~self._mask) | self._decode(n & self._mask)
+
+    def _decode(self, n):
+        result = 0
+        for i, b in enumerate(reversed(self._mapping)):
+            if n & (1 << b):
+                result |= (1 << i)
+        return result
+
+    def enbase(self, x):
+        result = self._enbase(x)
+        padding = self._alphabet[0] * (self._min_length - len(result))
+        return '%s%s' % (padding, result)
+
+    def _enbase(self, x):
+        n = len(self._alphabet)
+        if x < n:
+            return self._alphabet[x]
+        return self._enbase(int(x // n)) + self._alphabet[int(x % n)]
+
+    def debase(self, x):
+        n = len(self._alphabet)
+        result = 0
+        for i, c in enumerate(reversed(x)):
+            result += self._alphabet.index(c) * (n ** i)
+        return result
+
 class StateCorruptError(Exception):
     pass
 
+class URLMap:
+
+    _lock   = None
+
+    _dir_fd = None
+    _fname  = None
+    _file   = None
+
+    _next_id = None
+    _enc     = None
+
+    _full_to_short = None
+    _short_to_full = None
+
+    def __init__(self, state_dir_fd, fname, logger):
+        self._fname  = fname
+        self._dir_fd = os.dup(state_dir_fd)
+        self._enc    = UrlEncoder(block_size = 16)
+        self._lock   = threading.Lock()
+
+    def close(self):
+        if self._file is not None:
+            self._file.close()
+            self._file = None
+
+        if self._dir_fd is not None:
+            os.close(self._dir_fd)
+            self._dir_fd = None
+
+    def open(self):
+        if self._file is not None:
+            raise RuntimeError('Tried to open an already opened URL map')
+
+        # create the file if it does not exist
+        try:
+            with _open_dirfd(self._fname, 'w+x', perms = 0o600,
+                             dir_fd = self._dir_fd):
+                pass
+        except FileExistsError:
+            pass
+
+        try:
+            self._file = _open_dirfd(self._fname, 'r+', dir_fd = self._dir_fd)
+
+            data = [l.strip().split() for l in self._file.readlines()]
+            self._short_to_full = dict(data)
+            self._full_to_short = dict(((b, a) for (a, b) in data))
+            self._next_id       = self._enc.decode_url(data[-1][0]) if len(data) else 0
+        except:
+            self.close()
+            raise
+
+    def __enter__(self):
+        self.open()
+        return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+        self.close()
+
+    def short_to_full(self, short):
+        with self._lock:
+            return self._short_to_full[short]
+
+    def add(self, url):
+        with self._lock:
+            # mapping already exists, just return it
+            if url in self._full_to_short:
+                return self._full_to_short[url]
+
+            # find the next non-conflicting short id
+            short = None
+            for n in itertools.count(self._next_id):
+                short = self._enc.encode_url(n)
+                if not short in self._short_to_full:
+                    self._next_id = n + 1
+                    break
+
+            self._file.write('%s %s\n' % (short, url))
+            self._file.flush()
+
+            self._short_to_full[short] = url
+            self._full_to_short[url]   = short
+
+            return short
+
 class PersistentState:
     key = None
 
-    _fname = 'state'
+    urlmap = None
+
+    _fname_state = 'state'
+    _fname_map   = 'map'
 
-    def __init__(self, state_dir_fd, logger):
+
+    def __init__(self, state_dir_fd, public, logger):
         try:
-            with _open_dirfd(self._fname, 'w+x', perms = 0o600,
+            with _open_dirfd(self._fname_state, 'w+x', perms = 0o600,
                              dir_fd = state_dir_fd) as f:
                 data = { 'key' : secrets.token_hex(16) }
                 json.dump(data, f, indent = 4)
@@ -88,7 +236,7 @@ class PersistentState:
         except FileExistsError:
             pass
 
-        with _open_dirfd(self._fname, 'r', dir_fd = state_dir_fd) as f:
+        with _open_dirfd(self._fname_state, 'r', dir_fd = state_dir_fd) as f:
             try:
                 data = json.load(f)
             except json.decoder.JSONDecodeError:
@@ -97,6 +245,20 @@ class PersistentState:
 
         self.key = bytes.fromhex(data['key'])
 
+        if public:
+            self.urlmap = URLMap(state_dir_fd, self._fname_map, logger)
+
+    def close(self):
+        if self.urlmap:
+            self.urlmap.close()
+
+    def __enter__(self):
+        if self.urlmap:
+            self.urlmap.open()
+        return self
+    def __exit__(self, exc_type, exc_value, tb):
+        self.close()
+
 class HTTPChunkedRequestReader:
 
     _stream = None
@@ -201,6 +363,14 @@ class FShareRequestHandler(hs.BaseHTTPRequestHandler):
         # take the first path component, discard any extension
         fname = self._process_path(self.path).partition('/')[0]
         fname = os.path.splitext(fname)[0]
+        if self.server.state.urlmap:
+            try:
+                short = self.server.state.urlmap.short_to_full(fname)
+                self._logger.info('%s->%s', fname, short)
+                fname = short
+            except KeyError:
+                return self.send_error(HTTPStatus.NOT_FOUND)
+
         path = '/'.join((self.server.data_dir, fname))
         self._logger.info('serve file: %s', fname)
 
@@ -275,8 +445,16 @@ class FShareRequestHandler(hs.BaseHTTPRequestHandler):
         except KeyError:
             host = 'host.missing'
 
-        # the resulting URL is the secret HMAC + original basename
-        path  = urlparse.quote(dst_fname + '/' + src_fname)
+        if self.server.state.urlmap:
+            # public server: resulting URL is generated short URL + original extension
+            path  = self.server.state.urlmap.add(dst_fname)
+            self._logger.info('%s->%s', dst_fname, path)
+            path += os.path.splitext(src_fname)[1]
+        else:
+            # private srever: resulting URL is the secret HMAC + original basename
+            path  = dst_fname + '/' + src_fname
+
+        path  = urlparse.quote(path)
         reply = ('https://%s/%s' % (host, path)).encode('ascii')
 
         self.send_response(retcode)
@@ -345,6 +523,8 @@ parser = argparse.ArgumentParser('fshare server')
 
 parser.add_argument('-a', '--address', default = 'localhost')
 parser.add_argument('-p', '--port',    type = int, default = 5400)
+parser.add_argument('-P', '--public',  action = 'store_true',
+                    help = 'Generate public (short and guessable) URLs')
 
 group = parser.add_mutually_exclusive_group()
 group.add_argument('-4', '--ipv4',    action = 'store_true')
@@ -400,7 +580,7 @@ with contextlib.ExitStack() as stack:
 
     try:
         # read the state file
-        state = PersistentState(state_dir_fd, logger)
+        state = stack.enter_context(PersistentState(state_dir_fd, args.public, logger))
     except StateCorruptError:
         logger.error('Corrupted state file')
         sys.exit(1)