# -*- coding: utf-8 -*- """ Created on Sun Mar 22 10:46:20 2020 @author: cpan Simple client implementation with timeout feature to prevent hanging """ import sys import socket import json import io import struct import hashlib from Crypto import Random from Crypto.PublicKey import RSA HDRLEN = 2 class Message: def __init__(self, sock, addr): self._recv_buffer = bytearray(b"") self.sock = sock self.key = RSA.generate(2048, Random.new().read) self.sessionkey = None self.serverkey = None self.addr = addr def _json_encode(self,obj, encoding): return json.dumps(obj, ensure_ascii=False).encode(encoding) def _json_decode(self,json_bytes, encoding): tiow = io.TextIOWrapper( io.BytesIO(json_bytes), encoding=encoding, newline="" ) obj = json.load(tiow) tiow.close() return obj def _create_message(self, *, content_bytes, content_type, content_encoding): jsonheader = { "byteorder": sys.byteorder, "content-type": content_type, "content-encoding": content_encoding, "content-length": len(content_bytes), } jsonheader_bytes = self._json_encode(jsonheader, "utf-8") message_hdr = struct.pack(">H", len(jsonheader_bytes)) message = message_hdr + jsonheader_bytes + content_bytes return message def send_request(self, request): content = request["content"] content_type = request["type"] content_encoding = request["encoding"] if content_type == "text/json": req = { "content_bytes": self._json_encode(content, content_encoding), "content_type": content_type, "content_encoding": content_encoding, } else: req = { "content_bytes": content, "content_type": content_type, "content_encoding": content_encoding, } msg = self._create_message(**req) print(f"sending request {repr(msg)} to {self.addr}") self.sock.sendall(msg) def read_to_buffer(self, length): while len(self._recv_buffer) < length: self._recv_buffer += self.sock.recv(4096) def process_response(self): # read and process protoheader (fixed length of 2 byte in network order) self.read_to_buffer(HDRLEN) _jsonheader_len = struct.unpack(">H", self._recv_buffer[:HDRLEN])[0] self._recv_buffer = self._recv_buffer[HDRLEN:] # read and process jsonheader self.read_to_buffer(_jsonheader_len) jsonheader = self._json_decode(self._recv_buffer[:_jsonheader_len], "utf-8") self._recv_buffer = self._recv_buffer[_jsonheader_len:] for reqhdr in ("byteorder", "content-length", "content-type", "content-encoding",): if reqhdr not in jsonheader: raise ValueError(f'Missing required header "{reqhdr}".') content_len = jsonheader["content-length"] self.read_to_buffer(content_len) data = self._recv_buffer[:content_len] self._recv_buffer = self._recv_buffer[content_len:] if jsonheader["content-type"] == "text/json": encoding = jsonheader["content-encoding"] response = self._json_decode(data, encoding) print("received response", repr(response), "from", self.addr) content = response result = content.get("result") print(f"got result: {result}") else: # Binary or unknown content-type response = data print( f'received {jsonheader["content-type"]} response from', self.addr, ) content = response print(f"got response: {repr(content)}") def _send_handshake(self, data): response = { "content_bytes": data, "content_type": "binary/custom-server-binary-type", "content_encoding": "binary", } message = self._create_message(**response) self.sock.sendall(message) def close(self): self.sock.shutdown(socket.SHUT_RDWR) self.sock.close() self.sock = None def _recv_handshake(self): # read and process protoheader (fixed length of 2 byte in network order) self.read_to_buffer(HDRLEN) _jsonheader_len = struct.unpack(">H", self._recv_buffer[:HDRLEN])[0] self._recv_buffer = self._recv_buffer[HDRLEN:] # read and process jsonheader self.read_to_buffer(_jsonheader_len) jsonheader = self._json_decode(self._recv_buffer[:_jsonheader_len], "utf-8") self._recv_buffer = self._recv_buffer[_jsonheader_len:] for reqhdr in ("byteorder", "content-length", "content-type", "content-encoding",): if reqhdr not in jsonheader: raise ValueError(f'Missing required header "{reqhdr}".') content_len = jsonheader["content-length"] self.read_to_buffer(content_len) data = self._recv_buffer[:content_len] self._recv_buffer = self._recv_buffer[content_len:] sesskeyhash = data[-128:] data = data[:-128] encryptedsess = data[-256:] halfsesskey = self.key.decrypt(bytes(encryptedsess)) data = data[:-256] pubkeyhash = data[-128:] pubkey = data[:-128] if hashlib.sha3_512(pubkey).hexdigest().encode() != pubkeyhash: self.close() raise Exception('Server public key does not match hash') if hashlib.sha3_512(halfsesskey).hexdigest().encode() != sesskeyhash: self.close() raise Exception('Server public key does not match hash') self.sessionkey = halfsesskey self.serverkey = RSA.importKey(pubkey) print(f"Received server secrets correctly.") def send_passphrase(self, passphrase): passphrase = passphrase.encode() passhash = hashlib.sha3_512(passphrase).hexdigest().encode() sesshash = hashlib.sha3_512(self.sessionkey).hexdigest().encode() data = self.serverkey.encrypt(passhash + sesshash, None)[0] print(f'Sending passphrase + hash + sessionkey hash...') self._send_handshake(data) def create_request(action, value): if action == "search": return dict( type="text/json", encoding="utf-8", content=dict(action=action, value=value), ) else: return dict( type="binary/custom-client-binary-type", encoding="binary", content=bytes(action + value, encoding="utf-8"), ) def main(): if len(sys.argv) != 5: print("usage:", sys.argv[0], " ") sys.exit(1) host, port = sys.argv[1], int(sys.argv[2]) addr = (host, port) print("starting connection to", addr) sock = socket.create_connection(addr) sock.settimeout(None) message = Message(sock, addr) # security handshake publickey = message.key.publickey().exportKey('PEM') keyhash = hashlib.sha3_512(publickey).hexdigest().encode() data = publickey + keyhash print(f'Sending client public key and hash...') message._send_handshake(data) # receiving server secrets - public key and session key message._recv_handshake() message.send_passphrase('adminadmin') print('Handshake completed. Sever would close the sock if password wrong') action, value = sys.argv[3], sys.argv[4] request = create_request(action, value) message.send_request(request) # read and process response message.process_response() sock.shutdown(socket.SHUT_RDWR) sock.close() if __name__ == '__main__': main()