medashare/ui/server.py

#!/usr/bin/env python

# Notes:
# Python requests: https://2.python-requests.org/en/master/
# IRequest interface: https://twistedmatrix.com/documents/current/api/twisted.web.iweb.IRequest.html
# IResource interface: https://twistedmatrix.com/documents/current/api/twisted.web.resource.IResource.html
# Twisted TDD: https://twistedmatrix.com/documents/current/core/howto/trial.html
# Hypothesis: https://hypothesis.readthedocs.io/en/latest/
# Going Async from Flask to Twisted Klein: https://crossbario.com/blog/Going-Asynchronous-from-Flask-to-Twisted-Klein/
# Klein POST docs: https://klein.readthedocs.io/en/latest/examples/handlingpost.html

from klein import Klein
from kleintest import *
from twisted.trial import unittest
from twisted.web.iweb import IRequest
from cli import _asn1coder, Persona, MDBase, MetaData

import hashlib
import mock
import os.path
import shutil
import tempfile
import uuid

defaultfile = 'mediaserver.store.pasn1'
class MEDAServer:
	def __init__(self, fname):
		self._fname = fname
		self._hashes = {}
		try:
			data = _asn1coder.loads(open(fname).read())
			self._trustedkeys = {}
			for i in data['trustedkeys']:
				self.addpubkey(i)
			if data['objstore']:
				self._objstore = { uuid.UUID(bytes=k):
				    map(MDBase.create_obj, v) for k, v in
				    data['objstore'].iteritems() }
			else:
				self._objstore = {}
		except IOError:
			self._trustedkeys = {}
			self._objstore = {}

	app = Klein()

	def addpubkey(self, pubkey):
		persona = Persona.from_pubkey(pubkey)

		self._trustedkeys[persona.uuid] = persona

	def store(self):
		obj = {
			'trustedkeys': [ i.get_pubkey() for i in self._trustedkeys.itervalues() ],
			'objstore': self._objstore,
		}

		with open(self._fname, 'w') as fp:
			fp.write(_asn1coder.dumps(obj))

	@app.route('/lookup/<hash>')
	def lookup(self, request, hash):
		if hash in self._hashes:
			return

		request.setResponseCode(404)

	@app.route('/obj/<id>')
	def obj_lookup(self, request, id):
		try:
			id = uuid.UUID(id)
			return self._objstore[id][-1].encode()
		except ValueError:	# invalid format for uuid
			request.setResponseCode(400)
		except KeyError:	# no object
			request.setResponseCode(404)

	def _storeobj(self, obj):
		self._objstore.setdefault(obj.uuid, []).append(obj)
		try:
			hashes = obj.hashes
			for i in hashes:
				self._hashes.setdefault(i, []).append(obj.uuid)
		except AttributeError:
			pass

	@app.route('/store')
	def storeobj(self, request):
		try:
			obj = MDBase.decode(request.content.read())

			#if obj.type == 'identity':
			keyuuid = obj.uuid
			#else:
			#	keyuuid = obj.created_by_ref

			persona = self._trustedkeys[keyuuid]
			persona.verify(obj)

			self._storeobj(obj)

			request.setResponseCode(201)

		except Exception:
			request.setResponseCode(401)

# twistd support
#medaserver = MEDAServer()
#resource = medaserver.app.resource

def main():
	from optparse import OptionParser

	parser = OptionParser()
	parser.add_option('-a', action='append', dest='addpubkey',
	    default=[], help='Add specified public key as a trusted key.')

	options, args = parser.parse_args()

	medaserver = MEDAServer(defaultfile)

	try:
		if options.addpubkey:
			for i in options.addpubkey:
				medaserver.addpubkey(i)
			return

		medaserver.app.run()
	finally:
		medaserver.store()

if __name__ == '__main__':	# pragma: no cover
	main()

class _BaseServerTest(unittest.TestCase):
	def setUp(self):
		d = os.path.realpath(tempfile.mkdtemp())
		self.basetempdir = d
		self.medaserverfile = os.path.join(self.basetempdir, 'serverstore.pasn1')
		self.medaserver = MEDAServer(self.medaserverfile)
		self.requests = FakeRequests(self.medaserver.app)

	def tearDown(self):
		shutil.rmtree(self.basetempdir)
		self.basetempdir = None

class _TestCases(_BaseServerTest):
	def test_objlookup(self):
		# that when fetching an non-existant object
		r = self.requests.get('/obj/%s' % str(uuid.uuid4()))

		# it is 404, not found
		self.assertEqual(r.status_code, 404)

		# that passing an invalid uuid
		r = self.requests.get('/obj/bogusuuid')

		# it is 400, bad request
		self.assertEqual(r.status_code, 400)

	def test_pubkeystorage(self):
		# that an identity
		persona = Persona()
		persona.generate_key()

		# that by default, put's
		r = self.requests.put('/store', data=persona.get_identity().encode())

		# are denied
		self.assertEqual(r.status_code, 401)

		# can have it's public key added to the server
		self.medaserver.addpubkey(persona.get_pubkey())

		# that it can store the pubkey's identity
		r = self.requests.put('/store', data=persona.get_identity().encode())
		self.assertEqual(r.status_code, 201)

		# that an object with a bad signature
		badsigobj = persona.get_identity().new_version()

		# when stored
		r = self.requests.put('/store', data=badsigobj.encode())

		# is rejected
		self.assertEqual(r.status_code, 401)

		# that when fetching the object back
		r = self.requests.get('/obj/%s' % str(persona.uuid))

		# it is successful
		self.assertEqual(r.status_code, 200)

		# that the returned data
		fetchobj = MDBase.decode(r.text)

		# matches what was stored
		self.assertEqual(fetchobj, persona.get_identity())

		# that when stored
		self.medaserver.store()

		# and restarted
		self.medaserver = MEDAServer(self.medaserverfile)
		self.requests = FakeRequests(self.medaserver.app)

		# that fetching a previously successful object
		r = self.requests.get('/obj/%s' % str(persona.uuid))

		# it is successful
		self.assertEqual(r.status_code, 200)

	@mock.patch('klein.Klein.run')
	def test_addpubkey(self, apprun):
		persona = Persona()
		persona.generate_key()

		with mock.patch('server.MEDAServer.addpubkey') as addpub, \
		    mock.patch('sys.argv', [ 'progname', '-a',
		    persona.get_pubkey() ]) as argv:
			main()

		addpub.assert_called_with(persona.get_pubkey())
		apprun.assert_not_called()

	# Note: because of this mock, it hides the actual app.run call w/
	# a mock
	@mock.patch('server.MEDAServer')
	def test_medaserverinstanciated(self, medaserver):
		# that when main is run
		main()

		# that it gets called with the default storage file
		medaserver.assert_called_with('mediaserver.store.pasn1')

	@mock.patch('server.MEDAServer.store')
	@mock.patch('klein.Klein.run')
	def test_appruns(self, kleinrun, storefun):
		main()

		kleinrun.assert_called()
		storefun.assert_called()

class _TestPostConfig(_BaseServerTest):
	def setUp(self):
		_BaseServerTest.setUp(self)

		persona = Persona()
		persona.generate_key()

		self.persona = persona
		self.medaserver.addpubkey(persona.get_pubkey())

	def test_hashlookup(self):
		# that the hash of a file
		h = hashlib.sha256(open('fixtures/testfiles/test.txt', 'rb').read()).hexdigest()

		# when looked up
		print('hl')
		r = self.requests.get(('/lookup/%s' % h).encode('ascii'))

		# returns a 404
		self.assertEqual(r.status_code, 404)

		# but when the metadata object
		mdobj = self.persona.MetaData(hashes=[ h ], mimetype=[ 'text/plain' ])

		# is in the server
		self.medaserver._storeobj(mdobj)

		# when looked up
		r = self.requests.get('/lookup/%s' % h)

		# it is found
		self.assertEqual(r.status_code, 200)