from hashtree import HashTree

import collections

import socket

import logging as log

import config as conf

from networkers import NetworkReader,NetworkWriter

class Connection:

	def __init__(self):

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

		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.close()

class Client:

	def __init__(self,filename):

		self.filename=filename

	def negotiate(self):

		localTree=HashTree.fromFile(self.filename)

		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)

		# determine which blocks to send

		while len(nodeStack)>0:

			with Connection() as (incoming,outcoming):

				i=nodeStack.pop()

				outcoming.writeMsg({"command":"req", "index":i})

				jsonData,binData=incoming.readMsg()

				assert jsonData["index"]==i

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

				if localTree.store[i]!=binData:

					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

		return blocksToTransfer

	def sendData(self,blocksToTransfer):

		log.info(blocksToTransfer)

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

		i1=-1

		for i2 in blocksToTransfer:

			with Connection() as (incoming,outcoming):

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

				outcoming.writeMsg(jsonData,binData)

				i1=i2

		with Connection() as (incoming,outcoming):

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

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

		dataFile.close()

import hashlib

import os

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.index=self.leafStart

		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)

			for i in range(leafCount):

				data=f.read(HashTree.BLOCK_SIZE)

				res.insertLeaf(hashlib.sha256(data).digest()[HashTree.HASH_LEN:])

		res.buildTree()

		return res

	@classmethod

	def load(cls,filename):

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

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

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

		self.store[self.index]=h

		self.index+=1

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

			index=(index-1)//2

	## Fast construction of the tree over the leaves. O(n).

	def buildTree(self):

		for i in range(self.leafStart-1,-1,-1):

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

if __name__=="__main__":

	f1=HashTree.fromFile(open("serverFile.txt",mode='rb'))

	f2=HashTree.fromFile(open("clientFile.txt",mode='rb'))

	for i,(h1,h2) in enumerate(zip(f1.store,f2.store)):

		print("{0:2}".format(i),h1.hex(),h2.hex(),h1==h2)

import sys

import os.path

from argparse import ArgumentParser

from hashtree import HashTree

from client import Client

from server import Server

def _checkFile(f):

	if not os.path.isfile(f):

		print("invalid file specified:",f,file=sys.stderr)

		sys.exit(1)

def buildTree(args):

	_checkFile(args.datafile)

	tree=HashTree.fromFile(args.datafile)

	tree.save(args.treefile)

def update(args):

	_checkFile(args.datafile)

	c=Client(args.datafile)

	blocksToTransfer=c.negotiate()

	c.sendData(blocksToTransfer)

def serve(args):

	_checkFile(args.datafile)

	if args.tree:

		_checkFile(args.tree)

	# debug copy default file

	import shutil

	origFilename=args.datafile

	filename=origFilename+"_"

	shutil.copyfile(origFilename,filename)

	s=Server(filename,args.tree)

	s.serve()

parser=ArgumentParser()

subparsers=parser.add_subparsers()

pUpdate=subparsers.add_parser("update")

pUpdate.add_argument("-p","--port",type=int)

pUpdate.add_argument("--host",default="127.0.0.1")

pUpdate.add_argument("-t","--tree",help="stored hash tree location")

pUpdate.add_argument("datafile")

pUpdate.set_defaults(func=update)

pServe=subparsers.add_parser("serve")

pServe.add_argument("-p","--port",type=int)

pServe.add_argument("--host",default="")

pServe.add_argument("-t","--tree",help="stored hash tree location")

pServe.add_argument("datafile")

pServe.set_defaults(func=serve)

args=parser.parse_args()

args.func(args)