import collections

import socket

import ssl

import logging as log

from datetime import datetime

import config as conf

import stats

from util import Progress

from hashtree import HashTree

from networkers import NetworkReader,NetworkWriter

class Connection:

	def __init__(self):

		sock=socket.socket(socket.AF_INET, socket.SOCK_STREAM)

		sslContext=ssl.create_default_context(cafile=conf.peers)

		sslContext.check_hostname=False

		sslContext.load_cert_chain(conf.certfile,conf.keyfile)

		self._socket=sslContext.wrap_socket(sock)

		self._socket.connect((conf.hosts[0], conf.port))

		fr=self._socket.makefile(mode="rb")

		fw=self._socket.makefile(mode="wb")

		self.incoming=NetworkReader(fr)

		self.outcoming=NetworkWriter(fw)

	def __enter__(self):

		return self.incoming,self.outcoming

	def __exit__(self, exc_type, exc_val, exc_tb):

		self._socket.shutdown(socket.SHUT_RDWR)

		self._socket.close()

class Client:

	def __init__(self,filename):

		self._incoming=None

		self._outcoming=None

		self._filename=filename

		print(datetime.now(), "initializing...")

		self._localTree=HashTree.fromFile(self._filename)

	def negotiate(self):

		localTree=self._localTree

		blocksToTransfer=[]

		nodeStack=collections.deque([0]) # root

		# initialize session

		jsonData={"command":"init", "blockSize":localTree.BLOCK_SIZE, "blockCount":localTree.leafCount, "version":conf.version}

		self._outcoming.writeMsg(jsonData)

		jsonData,binData=self._incoming.readMsg()

		assert jsonData["command"]=="ack"

		# determine which blocks to send

		print(datetime.now(), "negotiating:")

		progress=Progress(localTree.leafCount)

		while len(nodeStack)>0:

			jsonData,binData=self._incoming.readMsg()

			assert jsonData["dataType"]=="hash"

		progress.done()

	def sendData(self,blocksToTransfer):

		log.info(blocksToTransfer)

		dataFile=open(self._filename, mode="rb")

		i1=-1

		print(datetime.now(), "sending data:")

		progress=Progress(len(blocksToTransfer))

		for (k,i2) in enumerate(blocksToTransfer):

			jsonData={"command":"send", "index":i2, "dataType":"data"}

			if i1+1!=i2:

				dataFile.seek(i2*HashTree.BLOCK_SIZE)

			binData=dataFile.read(HashTree.BLOCK_SIZE)

			log.info("block #{0}: {1}...{2}".format(i2,binData[:5],binData[-5:]))

			self._outcoming.writeMsg(jsonData,binData)

			stats.logTransferredBlock()

			jsonData,binData=self._incoming.readMsg()

			assert jsonData["command"]=="ack" and jsonData["index"]==i2, jsonData

			i1=i2

			progress.p(k)

		progress.done()

		self._outcoming.writeMsg({"command":"end"})

		log.info("closing session...")

		dataFile.close()

	def pullData(self,blocksToTransfer):

		log.info(blocksToTransfer)

		dataFile=open(self._filename, mode="rb+")

		i1=-1

		print(datetime.now(), "receiving data:")

		progress=Progress(len(blocksToTransfer))

		for (k,i2) in enumerate(blocksToTransfer):

			self._outcoming.writeMsg({"command":"req", "index":i2, "dataType":"data"})

			jsonData,binData=self._incoming.readMsg()

			assert jsonData["command"]=="send" and jsonData["index"]==i2 and jsonData["dataType"]=="data", jsonData

			if i1+1!=i2:

				dataFile.seek(i2*HashTree.BLOCK_SIZE)

			dataFile.write(binData)

			log.info("block #{0}: {1}...{2}".format(i2,binData[:5],binData[-5:]))

			stats.logTransferredBlock()

			i1=i2

			progress.p(k)

		progress.done()

		self._outcoming.writeMsg({"command":"end"})

		log.info("closing session...")

		dataFile.close()

	def setConnection(self,connection):

		(self._incoming,self._outcoming)=connection

import hashlib

import socket

import ssl

import multiprocessing

import logging as log

from hashtree import HashTree

from networkers import NetworkReader,NetworkWriter

import config as conf

class Connection:

	def __init__(self,serverSocket,sslContext):

		sock, address = serverSocket.accept()

		self._socket=sslContext.wrap_socket(sock,server_side=True)

		log.info('Connected by {0}'.format(address))

		fr=self._socket.makefile(mode="rb")

		fw=self._socket.makefile(mode="wb")

		self.incoming=NetworkReader(fr)

		self.outcoming=NetworkWriter(fw)

	def __enter__(self):

		return self.incoming,self.outcoming

	def __exit__(self, exc_type, exc_val, exc_tb):

		self._socket.shutdown(socket.SHUT_RDWR)

		self._socket.close()

class Miniserver:

	def __init__(self,filename,treeFile=""):

		self._filename=filename

		self._treeFile=treeFile

		self._ssl=ssl.create_default_context(ssl.Purpose.CLIENT_AUTH,cafile=conf.peers)

		self._ssl.verify_mode=ssl.CERT_REQUIRED

		self._ssl.load_cert_chain(conf.certfile,conf.keyfile)

		self._ss = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

		self._ss.bind(("", conf.port))

		self._ss.listen(1)

	def serve(self):

		while True:

			with Connection(self._ss,self._ssl) as c:

				p=multiprocessing.Process(target=Server.run,args=(c,self._filename,self._treeFile))

				p.start()

				p.join()

class Server:

	def __init__(self,connection,filename,treeFile=""):

		(self._incoming,self._outcoming)=connection

		self._filename=filename

		self._treeFile=treeFile

		if treeFile:

			self._tree=HashTree.load(treeFile)

		else:

			self._tree=HashTree.fromFile(filename)

		self._newLeaves=dict()

		self.BLOCK_SIZE=self._tree.BLOCK_SIZE

		self._lastIndex=-1

		self._dataFileHandle=None

	@staticmethod

	def run(*args):

		s=Server(*args)

		s.serve()

	@property

	def _dataFile(self):

		if not self._dataFileHandle:

			self._dataFileHandle=open(self._filename, mode="rb+")

		return self._dataFileHandle

	def serve(self):

		try:

			while self._serveOne(): pass

		except AssertionError as e:

			log.warning(e)

	def _serveOne(self):

		jsonData,binData=self._incoming.readMsg()

		if jsonData["command"]=="init":

			assert jsonData["blockSize"]==self.BLOCK_SIZE

			assert jsonData["blockCount"]==self._tree.leafCount

			self._outcoming.writeMsg({"command": "ack"})

		elif jsonData["command"]=="req":

			if jsonData["dataType"]=="data":

				self._outcoming.writeMsg(*self._requestData(jsonData["index"]))

			else:

				self._outcoming.writeMsg(*self._requestHash(jsonData["index"]))

		elif jsonData["command"]=="send" and jsonData["dataType"]=="data":

			self._outcoming.writeMsg(*self._receiveData(jsonData,binData))

		elif jsonData["command"]=="end":

			self._finalize()

			self._locked=False

			return False

		else:

			assert False, jsonData["command"]

		return True

		return (jsonResponse,binResponse)

	def _requestData(self,index):

		log.info("received request for data block #{0}".format(index))

		jsonResponse={"command":"send", "index":index, "dataType":"data"}

		if self._lastIndex+1!=index:

			self._dataFile.seek(index*self.BLOCK_SIZE)

		binResponse=self._dataFile.read(self.BLOCK_SIZE)

		return (jsonResponse,binResponse)

	def _receiveData(self,jsonData,binData):

		log.info("received data block #{0}: {1}...{2}".format(jsonData["index"],binData[:5],binData[-5:]))

		i=jsonData["index"]

		if self._lastIndex+1!=i:

			self._dataFile.seek(i*self.BLOCK_SIZE)

		self._dataFile.write(binData)

		self._lastIndex=i

		if self._treeFile:

			self._newLeaves[i+self._tree.leafStart]=hashlib.sha256(binData).digest()[HashTree.HASH_LEN:]

		return ({"command": "ack", "index": i},)

	def _finalize(self):

		log.info("closing session...")

		self._dataFile.close()

		self._dataFileHandle=None

		if self._treeFile:

			log.info("updating hash tree...")

			for (k,v) in self._newLeaves.items():

				self._tree.updateLeaf(k, v)

			self._tree.save(self._treeFile)

		log.info("done")

class Stats:

	received=0

	sent=0

	exchangedNodes=0

	transferredBlocks=0

def logReceived(data):

	Stats.received+=len(data)

def logSent(data):

	Stats.sent+=len(data)

def logTransferredBlock():

	Stats.transferredBlocks+=1

def report():

	return """received {r}B

sent {s}B

exchanged {nodes} hash tree nodes

transferred {blocks} blocks""".format(r=Stats.received, s=Stats.sent, nodes=Stats.exchangedNodes, blocks=Stats.transferredBlocks)