# -*- coding: utf-8 -*- """ Created on Sun Mar 22 10:44:50 2020 @author: cpan """ import sys import selectors import json import io import os import struct import hashlib from Crypto import Random from Crypto.PublicKey import RSA request_search = { "morpheus": "Follow the white rabbit. \U0001f430", "ring": "In the caves beneath the Misty Mountains. \U0001f48d", "\U0001f436": "\U0001f43e Playing ball! \U0001f3d0", } class Message: def __init__(self, selector, sock, addr, passhash): self.selector = selector self.sock = sock self.addr = addr self._recv_buffer = b"" self._send_buffer = b"" self._jsonheader_len = None self.jsonheader = None self.request = None self.response_created = False # added for security self.key = RSA.generate(2048, Random.new().read) self.sessionkey = os.urandom(8) self.passhash = passhash # pass phrase hash self.clientkey = None # client public key self.handshake = False print(f"initiating message associated with {addr}") def __del__(self): print(f"destructing messange associlated with {self.addr}") def _set_selector_events_mask(self, mode): """Set selector to listen for events: mode is 'r', 'w', or 'rw'.""" if mode == "r": events = selectors.EVENT_READ elif mode == "w": events = selectors.EVENT_WRITE elif mode == "rw": events = selectors.EVENT_READ | selectors.EVENT_WRITE else: raise ValueError(f"Invalid events mask mode {repr(mode)}.") self.selector.modify(self.sock, events, data=self) def _read(self): try: # Should be ready to read data = self.sock.recv(4096) except BlockingIOError: # Resource temporarily unavailable (errno EWOULDBLOCK) pass else: if data: self._recv_buffer += data else: raise RuntimeError("Peer closed.") def _write(self): if self._send_buffer: print("sending", repr(self._send_buffer), "to", self.addr) try: # Should be ready to write sent = self.sock.send(self._send_buffer) except BlockingIOError: # Resource temporarily unavailable (errno EWOULDBLOCK) pass else: self._send_buffer = self._send_buffer[sent:] # Close when the buffer is drained. The response has been sent. if sent and not self._send_buffer: self.close() 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 _create_response_json_content(self): action = self.request.get("action") if action == "search": query = self.request.get("value") answer = request_search.get(query) or f'No match for "{query}".' content = {"result": answer} else: content = {"result": f'Error: invalid action "{action}".'} content_encoding = "utf-8" response = { "content_bytes": self._json_encode(content, content_encoding), "content_type": "text/json", "content_encoding": content_encoding, } return response def _create_response_binary_content(self): response = { "content_bytes": b"First 10 bytes of request: " + self.request[:10], "content_type": "binary/custom-server-binary-type", "content_encoding": "binary", } return response def _create_handshake_content(self, data): response = { "content_bytes": data, "content_type": "binary/custom-server-binary-type", "content_encoding": "binary", } return response def process_events(self, mask): if mask & selectors.EVENT_READ: self.read() if mask & selectors.EVENT_WRITE: self.write() def read(self): self._read() if self._jsonheader_len is None: self.process_protoheader() if self._jsonheader_len is not None: if self.jsonheader is None: self.process_jsonheader() if self.jsonheader: if self.request is None: self.process_request() def write(self): if self.request: if not self.response_created: self.create_response() self._write() def close(self): print(f"closing connection to {self.addr}\n") try: self.selector.unregister(self.sock) except Exception as e: print( f"error: selector.unregister() exception for", f"{self.addr}: {repr(e)}", ) try: self.sock.close() except OSError as e: print( f"error: socket.close() exception for", f"{self.addr}: {repr(e)}", ) finally: # Delete reference to socket object for garbage collection self.sock = None def process_protoheader(self): hdrlen = 2 if len(self._recv_buffer) >= hdrlen: self._jsonheader_len = struct.unpack( ">H", self._recv_buffer[:hdrlen] )[0] self._recv_buffer = self._recv_buffer[hdrlen:] def process_jsonheader(self): hdrlen = self._jsonheader_len if len(self._recv_buffer) >= hdrlen: self.jsonheader = self._json_decode( self._recv_buffer[:hdrlen], "utf-8" ) self._recv_buffer = self._recv_buffer[hdrlen:] for reqhdr in ( "byteorder", "content-length", "content-type", "content-encoding", ): if reqhdr not in self.jsonheader: raise ValueError(f'Missing required header "{reqhdr}".') def send_my_secrets(self): halfsesskey = self.sessionkey encryptedsess = self.clientkey.encrypt(halfsesskey, None)[0] sesskeyhash = hashlib.sha3_512(halfsesskey).hexdigest().encode() mypubkey = self.key.publickey().exportKey('PEM') mypubkeyhash = hashlib.sha3_512(mypubkey).hexdigest().encode() mysecret = mypubkey + mypubkeyhash + encryptedsess + sesskeyhash response = self._create_handshake_content(mysecret) message = self._create_message(**response) self._send_buffer += message print(f'Sending server public key, hash, session key, hash...') def process_request(self): content_len = self.jsonheader["content-length"] if not len(self._recv_buffer) >= content_len: return data = self._recv_buffer[:content_len] self._recv_buffer = self._recv_buffer[content_len:] if not self.handshake: # data is client public key plus hash key = data[:-128] keyhash = data[-128:] if hashlib.sha3_512(key).hexdigest().encode() == keyhash: self.clientkey = RSA.importKey(key) else: self.sock.close() raise Exception('Client plublic key does not match key hash') self.send_my_secrets() else: if self.jsonheader["content-type"] == "text/json": encoding = self.jsonheader["content-encoding"] self.request = self._json_decode(data, encoding) print("received request", repr(self.request), "from", self.addr) else: # Binary or unknown content-type self.request = data print( f'received {self.jsonheader["content-type"]} request from', self.addr, ) # Set selector to listen for write events, we're done reading. self._set_selector_events_mask("w") def create_response(self): if self.jsonheader["content-type"] == "text/json": response = self._create_response_json_content() else: # Binary or unknown content-type response = self._create_response_binary_content() message = self._create_message(**response) self.response_created = True self._send_buffer += message