diff --git a/src/client.py b/src/client.py --- a/src/client.py +++ b/src/client.py @@ -1,5 +1,6 @@ import collections import socket +import ssl import logging as log from datetime import datetime @@ -11,8 +12,10 @@ from networkers import NetworkReader,Net class Connection: - def __init__(self): - self._socket=socket.socket(socket.AF_INET, socket.SOCK_STREAM) + def __init__(self,sslContext): + sock=socket.socket(socket.AF_INET, socket.SOCK_STREAM) + + 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") @@ -24,6 +27,7 @@ class Connection: return self.incoming,self.outcoming def __exit__(self, exc_type, exc_val, exc_tb): + self._socket.shutdown(socket.SHUT_RDWR) self._socket.close() @@ -31,6 +35,10 @@ class Client: def __init__(self,filename): self._filename=filename + self._ssl=ssl.create_default_context(cafile=conf.peers) + self._ssl.check_hostname=False + self._ssl.load_cert_chain(conf.certfile,conf.keyfile) + def negotiate(self): print(datetime.now(), "initializing...") localTree=HashTree.fromFile(self._filename) @@ -38,7 +46,7 @@ class Client: nodeStack=collections.deque([0]) # root # initialize session - with Connection() as (incoming,outcoming): + with Connection(self._ssl) as (incoming,outcoming): jsonData={"command":"init", "blockSize":localTree.BLOCK_SIZE, "blockCount":localTree.leafCount, "version":conf.version} outcoming.writeMsg(jsonData) jsonData,binData=incoming.readMsg() @@ -49,7 +57,7 @@ class Client: progress=Progress(localTree.leafCount) while len(nodeStack)>0: i=nodeStack.pop() - outcoming.writeMsg({"command":"req", "index":i}) + outcoming.writeMsg({"command":"req", "index":i, "dataType":"hash"}) jsonData,binData=incoming.readMsg() assert jsonData["index"]==i @@ -73,7 +81,7 @@ class Client: i1=-1 print(datetime.now(), "sending data:") - with Connection() as (incoming,outcoming): + with Connection(self._ssl) as (incoming,outcoming): progress=Progress(len(blocksToTransfer)) for (k,i2) in enumerate(blocksToTransfer): jsonData={"command":"send", "index":i2, "dataType":"data"} @@ -91,7 +99,7 @@ class Client: progress.p(k) progress.done() - with Connection() as (incoming,outcoming): + with Connection(self._ssl) as (incoming,outcoming): outcoming.writeMsg({"command":"end"}) log.info("closing session...")