from hashtree import HashTree

import collections

import socket

import sys

import logging as log

import config as conf

from networkers import NetworkReader,NetworkWriter

filename=sys.argv[1]

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

def negotiate():

	localTree=HashTree.fromFile(open(filename,mode="rb"))

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

	log.info(blocksToTransfer)

	dataFile=open(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()

if __name__=="__main__":

	blocksToTransfer=negotiate()

	sendData(blocksToTransfer)

	sys.exit(0)

import hashlib

import os

import collections

class HashTree:

	HASH_LEN=16 # bytes

	BLOCK_SIZE=4096 # bytes

	## Prepares a tree containing leafCount leaves.

	def __init__(self,leafCount):

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

		self.leafStart=leafCount-1

		self.index=self.leafStart

		self.leafCount=leafCount

	@classmethod

	def fromFile(cls,fd):

		stat=os.fstat(fd.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=fd.read(HashTree.BLOCK_SIZE)

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

		res.buildTree()

		return res

	## 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 json

class NetworkReader:

	def __init__(self,stream):

		self.stream=stream

	def readMsg(self):

		data=self.stream.readline()

		if not data: pass # !! raise something

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

		data=self.stream.readline()

		if not data: pass # !! raise something

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

		jsonData=self.stream.read(jsonLength)

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

		binData=self.stream.read(binLength)

		return (jsonData,binData)

class NetworkWriter:

	def __init__(self,stream):

		self.stream=stream

	def writeMsg(self,*args):

		self.stream.write(self.prepMsg(*args))

		self.stream.flush()

	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 socket

from hashtree import HashTree

from networkers import NetworkReader,NetworkWriter

import collections

import sys

import logging as log

import config as conf

# debug copy default file

import shutil

origFilename=sys.argv[1]

filename=origFilename+"_"

shutil.copyfile(origFilename,filename)

class Connection:

	def __init__(self,server_socket):

		self.socket, address = server_socket.accept()

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

localTree=HashTree.fromFile(open(filename,mode="rb"))

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

ss.bind(("",conf.port))

ss.listen(1)

blocksToTransfer=[]

nodeStack=collections.deque([0])

i1=-1

while True:

	with Connection(ss) as (incoming,outcoming):

		jsonData,binData=incoming.readMsg()

		dataFile=open(filename,mode="rb+")

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

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

			assert jsonData["blockCount"]==localTree.leafCount

		elif jsonData["command"]=="req": # !! index out of range

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

			nodeHash=localTree.store[jsonData["index"]]

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

			binResponse=nodeHash

			outcoming.writeMsg(jsonResponse,binResponse)

		elif jsonData["command"]=="send" and jsonData["dataType"]=="data": # needlessly allow hashes and data in mixed order

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

			i2=jsonData["index"]

			if i1+1!=i2:

				dataFile.seek(i2*localTree.BLOCK_SIZE)

			dataFile.write(binData)

			i1=i2

			# never update the hash tree

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

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

			break

		else: pass # !! error

dataFile.close()

sys.exit(0)