async def initial_connection_check(self, reader: StreamReader, writer: StreamWriter): host_port = writer.get_extra_info('peername') new_user: Optional[User] = None try: # 1. send plain message writer.write(b'hello') await writer.drain() # 2. receive other's header try: received = await asyncio.wait_for(reader.read(BUFFER_SIZE), 5.0) if len(received) == 0: raise PeerToPeerError('empty msg receive') header = json.loads(received.decode()) except asyncio.TimeoutError: raise PeerToPeerError('timeout on other\'s header receive') except json.JSONDecodeError: raise PeerToPeerError( 'json decode error on other\'s header receive') # 3. generate new user user_header = UserHeader(**header) new_user = User(user_header, self.number, reader, writer, host_port, AESCipher.create_key(), SERVER_SIDE) self.number += 1 if new_user.header.name == V.SERVER_NAME: raise ConnectionAbortedError('Same origin connection') # 4. send my public key my_sec, my_pub = generate_keypair() send = json.dumps({'public-key': my_pub}).encode() await new_user.send(send) self.traffic.put_traffic_up(send) # 5. receive public key try: receive = await new_user.recv() self.traffic.put_traffic_down(receive) if len(receive) == 0: raise ConnectionAbortedError('received msg is zero.') data = json.loads(receive.decode()) except asyncio.TimeoutError: raise PeerToPeerError('timeout on public key receive') except json.JSONDecodeError: raise PeerToPeerError( 'json decode error on public key receive') # 6. encrypt and send AES key and header send = json.dumps({ 'aes-key': new_user.aeskey, 'header': self.get_server_header(), }) key = generate_shared_key(my_sec, data['public-key']) encrypted = AESCipher.encrypt(key, send.encode()) await new_user.send(encrypted) self.traffic.put_traffic_up(encrypted) # 7. receive accept signal try: encrypted = await new_user.recv() self.traffic.put_traffic_down(encrypted) except asyncio.TimeoutError: raise PeerToPeerError('timeout on accept signal receive') receive = AESCipher.decrypt(new_user.aeskey, encrypted) if receive != b'accept': raise PeerToPeerError(f"Not accept signal! {receive}") # 8. accept connection log.info( f"established connection as server from {new_user.header.name} {new_user.get_host_port()}" ) asyncio.ensure_future(self.receive_loop(new_user)) # server port's reachable check await asyncio.sleep(1.0) await self.check_reachable(new_user) return except (ConnectionAbortedError, ConnectionResetError) as e: msg = f"disconnect error {host_port} {e}" except PeerToPeerError as e: msg = f"peer2peer error {host_port} {e}" except Exception as e: msg = "InitialConnCheck: {}".format(e) log.error(msg, exc_info=True) # EXCEPTION! if new_user: # remove user self.remove_connection(new_user, msg) else: # close socket log.debug(msg) try: writer.write(msg.encode()) await writer.drain() except Exception: pass try: writer.close() except Exception: pass
def _initial_connection_check(self, sock, host_port): sock.settimeout(10) try: # ヘッダーを受取る received = sock.recv(self.buffsize) if len(received) == 0: raise ConnectionAbortedError('zero msg, connection closed.') self.traffic.put_traffic_down(received) header = json.loads(received.decode()) with self.lock: new_user = User(self.number, sock, host_port, aeskey=AESCipher.create_key(), sock_type=C.T_SERVER) self.number += 1 new_user.deserialize(header) if new_user.name == V.SERVER_NAME: raise ConnectionAbortedError('Same origin connection.') # こちらの公開鍵を送る send = json.dumps({'public-key': self.ecc.pk}).encode() sock.sendall(send) self.traffic.put_traffic_up(send) # 公開鍵を取得する receive = new_user.sock.recv(self.buffsize) self.traffic.put_traffic_down(receive) if len(receive) == 0: raise ConnectionAbortedError('received msg is zero.') public_key = json.loads(receive.decode())['public-key'] # AESKEYとHeaderを暗号化して送る encrypted = self.ecc.encrypt(recipient_pk=public_key, msg=json.dumps({ 'aes-key': new_user.aeskey, 'header': self.get_server_header() }).encode(), encode='raw') new_user.send(encrypted) self.traffic.put_traffic_up(encrypted) # Accept信号を受け取る encrypted = new_user.sock.recv(self.buffsize) self.traffic.put_traffic_down(encrypted) receive = AESCipher.decrypt(new_user.aeskey, encrypted) if receive != b'accept': raise ConnectionAbortedError('Not accept signal.') # Accept connection logging.info("New connection from \"{}\" {}".format( new_user.name, new_user.get_host_port())) Thread(target=self._receive_msg, name='S:' + new_user.name, args=(new_user, ), daemon=True).start() # Port accept check sleep(10) if new_user in self.user: self.is_reachable(new_user) return except ConnectionAbortedError as e: error = "ConnectionAbortedError, {}".format(e) except json.decoder.JSONDecodeError as e: error = "JSONDecodeError, {}".format(e) except socket.timeout: error = "socket.timeout" except Exception as e: error = "Exception as {}".format(e) # close socket error = "Close on initial check " + error logging.debug(error) try: sock.sendall(error.encode()) except: pass try: sock.close() except: pass