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):

		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):

class Client:

	def __init__(self,filename):

	def negotiate(self):

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

		blocksToTransfer=[]

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

		# initialize session

		with Connection() as (incoming,outcoming):

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

			outcoming.writeMsg(jsonData)

			jsonData,binData=incoming.readMsg()

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

			# determine which blocks to send

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

					if 2*i+3<len(localTree.store): # inner node

						nodeStack.append(2*i+2)

						nodeStack.append(2*i+1)

					else:

						blocksToTransfer.append(i-localTree.leafStart) # leaf

						progress.p(i-localTree.leafStart)

		progress.done()

		return blocksToTransfer

	def sendData(self,blocksToTransfer):

		log.info(blocksToTransfer)

		i1=-1

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

		with Connection() as (incoming,outcoming):

			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:]))

from util import Progress

class HashTree:

	HASH_LEN=16 # bytes

	BLOCK_SIZE=4096 # bytes

	## Prepares a tree containing leafCount leaves.

	def __init__(self,leafCount):

		self.store=[b""]*(leafCount*2-1)

		self.leafStart=leafCount-1

		self.leafCount=leafCount

	@classmethod

	def fromFile(cls,filename):

		with open(filename,"rb") as f:

			stat=os.fstat(f.fileno())

			size=stat.st_size # !! symlinks

			leafCount=(size-1)//HashTree.BLOCK_SIZE+1 # number of leaf blocks

			res=cls(leafCount)

			print(datetime.now(), "hashing file:")

			progress=Progress(leafCount)

			for i in range(leafCount):

			size=stat.st_size

			nodeCount=size//HashTree.HASH_LEN

			res=cls((nodeCount+1)//2)

			for i in range(nodeCount):

				res.store[i]=f.read(HashTree.HASH_LEN)

		return res

	def save(self,filename):

		with open(filename,"wb") as f:

			for h in self.store:

				f.write(h)

	## Inserts a leaf at the first empty position.

	#	Useful and used only during the tree construction.

	def insertLeaf(self,h):

	## Updates a hash stored in the leaf.

	def updateLeaf(self,index,h):

		if index<self.leafStart: raise IndexError()

		self.store[index]=h

		self.updateNode((index-1)//2)

	## Updates the node at index and all its ancestors.

	def updateNode(self,index):

		while index>=0:

			self.store[index]=hashlib.sha256(self.store[index*2+1]+self.store[index*2+2]).digest()[HashTree.HASH_LEN:]

import json

import stats

class NetworkReader:

	def __init__(self,stream):

	def readMsg(self):

		assert data

		stats.logReceived(data)

		jsonLength=int(data.split(b":")[1].strip()) # "json-length: length" -> length

		assert data

		stats.logReceived(data)

		binLength=int(data.split(b":")[1].strip()) # "bin-length: length" -> length

		assert len(jsonData)==jsonLength

		stats.logReceived(jsonData)

		jsonData=json.loads(str(jsonData,encoding="utf-8"))

		assert len(binData)==binLength

		stats.logReceived(binData)

		return (jsonData,binData)

class NetworkWriter:

	def __init__(self,stream):

	def writeMsg(self,*args):

		msg=self.prepMsg(*args)

		stats.logSent(msg)

	def prepMsg(self,jsonData,binData=b""):

		jsonData=bytes(json.dumps(jsonData)+"\n",encoding="utf-8")

		jsonLength=bytes("json-length: "+str(len(jsonData))+"\n",encoding="utf-8")

		binLength=bytes("bin-length: "+str(len(binData))+"\n",encoding="utf-8")

		return b"".join((jsonLength,binLength,jsonData,binData))

import hashlib

import socket

import logging as log

from hashtree import HashTree

from networkers import NetworkReader,NetworkWriter

import config as conf

class Connection:

	def __init__(self,serverSocket):

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

		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):

class Server:

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

		self._treeFile=treeFile

		if treeFile:

		else:

		self._newLeaves=dict()

		self._lastWrite=-1

	def serve(self):

		while True:

				try:

					while True:

						if not self._serveOne(incoming,outcoming): return

				except AssertionError:

					continue

	def _serveOne(self,incoming,outcoming):

		jsonData,binData=incoming.readMsg()

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

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

			outcoming.writeMsg({"command": "ack"})

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

			outcoming.writeMsg(*self._requestHash(jsonData))

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

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

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

			self._finalize()

			return False

		else:

			assert False, jsonData["command"]

		return True

	def _requestHash(self,jsonData):

		log.info("received request for node #{0}".format(jsonData["index"]))

		jsonResponse={"command":"send", "index":jsonData["index"], "dataType":"hash"}

		binResponse=nodeHash

		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._lastWrite+1!=i:

		self._lastWrite=i

		if self._treeFile:

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

	def _finalize(self):

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

		if self._treeFile:

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

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

		log.info("done")