Exemplo n.º 1
0
    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
Exemplo n.º 2
0
 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