memcache_repl.py

#!/usr/bin/env python
from __future__ import unicode_literals
import argparse
import cmd
import hashlib
import shlex
import socket

try:
    import readline  # noqa
except ImportError:
    pass  # no real harm, just less usable


class HashingMemcacheClient(cmd.Cmd, object):
    prompt = 'memcache> '
    hashfunc = hashlib.md5

    def __init__(self, host, port):
        super(HashingMemcacheClient, self).__init__()
        self._socket = socket.socket()
        self._socket.connect((host, port))

    def send(self, message):
        self._socket.sendall(message + b'\r\n')
        # The responses should all be pretty tiny.
        return self._socket.recv(1024).decode('utf-8')

    def hash(self, message):
        hash = self.hashfunc(message.encode('utf-8')).hexdigest()
        return hash.encode('ascii')

    def do_get(self, line):
        key, = self.parse(line)
        self.stdout.write(self.send(b'get ' + self.hash(key)))

    def do_delete(self, line):
        key, = self.parse(line)
        self.stdout.write(self.send(b'delete ' + self.hash(key)))

    def do_set(self, line):
        key, flags, ttl, value = self.parse(line)
        flags = int(flags)
        ttl = int(ttl)
        value = value.encode('utf-8')
        self.stdout.write(self.send(b'set %s %d %d %d\r\n%s' % (
            self.hash(key), flags, ttl, len(value), value)))

    def do_EOF(self, ignored=None):
        self._socket.close()
        return True

    def parse(self, line):
        return shlex.split(line)

    def onecmd(self, line):
        try:
            return super(HashingMemcacheClient, self).onecmd(line)
        except Exception:
            self.do_EOF()
            raise


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('host')
    parser.add_argument('port', type=int, default=11211, nargs='?')
    parsed_args = parser.parse_args()
    HashingMemcacheClient(parsed_args.host, parsed_args.port).cmdloop()