
import time
import tempfile
import os


class KeyWriter(object):
    """Generate and automatically expire random keys.

    The keys are stored in a filesystem directory.
    """

    def __init__(self, dirpath, length=48, freshness=300,
            max_age=3600, max_future=300):
        self.dirpath = dirpath
        self.length = length
        self.freshness = freshness
        self.max_age = max_age
        self.max_future = max_future
        self.current_key_mtime = 0
        self.current_key_id = None
        self.current_key = None

    def get_fresh_key(self):
        """Return a fresh (key_id, key) tuple."""
        now = time.time()
        if (not self.current_key_mtime
                or now >= self.current_key_mtime + self.freshness):
            self._prune()
            if not self._find_fresh_key():
                self._generate_key(now)
        return self.current_key_id, self.current_key

    def _generate_key(self, now):
        f = tempfile.NamedTemporaryFile(
            prefix='', dir=self.dirpath, delete=False)
        key_id = os.path.basename(f.name).encode('ascii')
        key = os.urandom(self.length)
        f.write(key)
        f.close()
        self.current_key_id = key_id
        self.current_key = key
        self.current_key_mtime = now

    def _list_keys(self):
        for name in os.listdir(self.dirpath):
            if not name.startswith('.'):
                fn = os.path.join(self.dirpath, name)
                if os.path.isfile(fn):
                    mtime = os.path.getmtime(fn)
                    yield name, fn, mtime

    def _find_fresh_key(self):
        """Look for a fresh key produced by another KeyWriter instance."""
        now = time.time()
        min_mtime = now - self.freshness / 2
        max_mtime = now + self.max_future
        for name, fn, mtime in self._list_keys():
            if mtime > min_mtime and mtime <= max_mtime:
                f = open(fn, 'rb')
                key = f.read()
                f.close()
                if len(key) == self.length:
                    # Use this key.
                    self.current_key_id = name.encode('ascii')
                    self.current_key = key
                    self.current_key_mtime = mtime
                    return True
        return False

    def _prune(self):
        now = time.time()
        min_mtime = now - self.max_age
        max_mtime = now + self.max_future
        for _name, fn, mtime in self._list_keys():
            if mtime <= min_mtime or mtime > max_mtime:
                # Too old or too new.
                os.remove(fn)


class KeyReader(object):
    """Read keys generated by a KeyWriter.

    Ignore expired keys.
    """

    def __init__(self, dirpath, length=48, max_age=3600, max_future=300):
        self.dirpath = dirpath
        self.length = length
        self.max_age = max_age
        self.max_future = max_future
        self.keys = {}  # bytes(key_id): (create_time, key, filename)

    def get_key(self, key_id):
        """Get the specified key.

        Raise KeyError if the key does not exist, is too old, is too new,
        or has the wrong length.
        """
        if not isinstance(key_id, bytes):
            raise TypeError("key_id must be a bytes object")
        now = time.time()
        try:
            mtime, key, fn = self.keys[key_id]
        except KeyError:
            if (not key_id
                    or key_id.startswith(b'.')
                    or b'/' in key_id
                    or b'\\' in key_id):
                raise
            fn = os.path.join(self.dirpath, key_id.decode('ascii'))
            if not os.path.isfile(fn):
                raise
            mtime = os.path.getmtime(fn)
            min_mtime = now - self.max_age
            max_mtime = now + self.max_future
            if mtime <= min_mtime or mtime > max_mtime:
                # The file is too old or too new.
                raise
            f = open(fn)
            key = f.read()
            f.close()
            if len(key) != self.length:
                # Wrong key length.
                raise
            self.keys[key_id] = (mtime, key, fn)
        else:
            if now >= mtime + self.max_age:
                # The key in self.keys is too old.
                self._prune()
                raise KeyError(key_id)
            if not os.path.isfile(fn):
                # The key has been removed.
                self._prune()
                raise KeyError(key_id)
        return key

    def _prune(self):
        now = time.time()
        min_mtime = now - self.max_age
        for key_id, (mtime, _key, fn) in self.keys.items():
            if mtime <= min_mtime:
                self.keys.pop(key_id, None)
            if not os.path.isfile(fn):
                self.keys.pop(key_id, None)
