diff --git a/src/server.py b/src/server.py --- a/src/server.py +++ b/src/server.py @@ -1,8 +1,9 @@ -import socket +import hashlib +import socket +import logging as log + from hashtree import HashTree from networkers import NetworkReader,NetworkWriter -import logging as log - import config as conf @@ -26,12 +27,14 @@ class Connection: class Server: def __init__(self,filename,treeFile=""): 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.ss = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -65,9 +68,7 @@ class Server: outcoming.writeMsg(*self._receiveData(jsonData,binData)) elif jsonData["command"]=="end": - log.info("closing session...") - if self.dataFile: - self.dataFile.close() + self._finalize() return False else: @@ -95,6 +96,18 @@ class Server: self.dataFile.seek(i*self.BLOCK_SIZE) self.dataFile.write(binData) self._lastWrite=i + if self._treeFile: + self._newLeaves[i+self.tree.leafStart]=hashlib.sha256(binData).digest()[HashTree.HASH_LEN:] return ({"command": "ack", "index": i},) - # never update the hash tree + + def _finalize(self): + log.info("closing session...") + if self.dataFile: + self.dataFile.close() + if self._treeFile: + log.info("updating hash tree...") + for (k,v) in self._newLeaves: + self.tree.updateLeaf(k,v) + self.tree.save(self._treeFile) + log.info("done")