
    i^                       d Z ddlmZ ddlZddlZddlZddlmZmZ ddl	m
Z
 ddlmZmZ erddlmZ dd	lmZ d
  ed          D             Z G d d          Z G d d          ZdZdZdZdZdZeZdZdZdZdZdZ dZ!dZ" G d de
          Z# G d de
          Z$e$j%        e$j&        e$j'        fZ(dZ)dZ*dZ+d Z,d!Z-d"Z.d#Z/d#Z0 ej1        d$d          Z2 G d% d&e3          Z4 G d' d(e          Z5 G d) d*e          Z6 G d+ d,e          Z7d;d2Z8 G d3 d4          Z9 G d5 d6          Z: G d7 d8          Z; G d9 d:          Z<dS )<zY
wsproto/frame_protocol
~~~~~~~~~~~~~~~~~~~~~~

WebSocket frame protocol implementation.
    )annotationsN)IncrementalDecodergetincrementaldecoder)IntEnum)TYPE_CHECKING
NamedTuple)	Generator   )	Extensionc                `    g | ]*t          fd t          d          D                       +S )c              3  "   K   | ]	}|z  V  
d S N ).0abs     A/usr/local/lib/python3.11/dist-packages/wsproto/frame_protocol.py	<genexpr>z<listcomp>.<genexpr>   s'      ..aAE......       )bytesrange)r   r   s    @r   
<listcomp>r      s<    CCC1e....5::.....CCCr   r   c                      e Zd Zd
dZddZd	S )XorMaskerSimplemasking_keybytearray | bytesreturnNonec                    || _         d S r   )_masking_key)selfr   s     r   __init__zXorMaskerSimple.__init__   s    'r   data	bytearrayc                   t          |          }|r|}d | j        D             \  }}}}|d d d                             |          |d d d<   |dd d                             |          |dd d<   |dd d                             |          |dd d<   |dd d                             |          |dd d<   t          |          dz  }| j        |d          | j        d |         z   | _        |S |S )Nc              3  0   K   | ]}t           |         V  d S r   )
_XOR_TABLE)r   ns     r   r   z*XorMaskerSimple.process.<locals>.<genexpr>!   s&      CCA*Q-CCCCCCr      r
         )r%   r!   	translatelen)r"   r$   
data_arrayr   r   cdkey_rotations           r   processzXorMaskerSimple.process   s     	JCC1BCCCJAq!Q(1o77::JsssO)!$Q$/99!<<Jqt!t)!$Q$/99!<<Jqt!t)!$Q$/99!<<Jqt!t t99q=L!,--043D]l]3SS  r   N)r   r   r   r   r$   r%   r   r%   )__name__
__module____qualname__r#   r3   r   r   r   r   r      s<        ( ( ( (     r   r   c                      e Zd ZddZdS )XorMaskerNullr$   r%   r   c                    |S r   r   r"   r$   s     r   r3   zXorMaskerNull.process3   s    r   Nr4   )r5   r6   r7   r3   r   r   r   r9   r9   2   s(             r   r9   ~      }   i  l       @             c                  2    e Zd ZdZdZdZdZdZdZdZ	dd
Z
dS )Opcodez7
    RFC 6455, Section 5.2 - Base Framing Protocol
    r   r
   r+      	   
   r   boolc                &    t          | dz            S )NrF   )rI   r"   s    r   	iscontrolzOpcode.iscontrold   s    D4K   r   Nr   rI   )r5   r6   r7   __doc__CONTINUATIONTEXTBINARYCLOSEPINGPONGrL   r   r   r   rE   rE   M   s\         
 L D F E D D! ! ! ! ! !r   rE   c                  J    e Zd ZdZdZdZdZdZdZdZ	dZ
d	Zd
ZdZdZdZdZdZdS )CloseReasonz8
    RFC 6455, Section 7.4.1 - Defined Status Codes
      i  i  i  i  i  i  i  i  i  i  i  i  i  N)r5   r6   r7   rN   NORMAL_CLOSURE
GOING_AWAYPROTOCOL_ERRORUNSUPPORTED_DATANO_STATUS_RCVDABNORMAL_CLOSUREINVALID_FRAME_PAYLOAD_DATAPOLICY_VIOLATIONMESSAGE_TOO_BIGMANDATORY_EXTINTERNAL_ERRORSERVICE_RESTARTTRY_AGAIN_LATERTLS_HANDSHAKE_FAILEDr   r   r   rV   rV   h   s          N J N  N  "& 
 O M
 N O O  r   rV   rW   i  i  i  i  i  z!Ic                  .     e Zd Zej        fd fdZ xZS )	ParseFailedmsgstrcoderV   r   r   c                X    t                                          |           || _        d S r   )superr#   rj   )r"   rh   rj   	__class__s      r   r#   zParseFailed.__init__   s(     				r   )rh   ri   rj   rV   r   r   )r5   r6   r7   rV   rZ   r#   __classcell__)rm   s   @r   rg   rg      sH        ,7,F          r   rg   c                  .    e Zd ZU ded<   ded<   ded<   dS )RsvBitsrI   rsv1rsv2rsv3Nr5   r6   r7   __annotations__r   r   r   rp   rp      s+         JJJJJJJJJJJr   rp   c                  B    e Zd ZU ded<   ded<   ded<   ded<   d	ed
<   dS )HeaderrI   finrp   rsvrE   opcodeintpayload_lenbytes | Noner   Nrt   r   r   r   rw   rw      sE         IIILLLNNNr   rw   c                  8    e Zd ZU ded<   ded<   ded<   ded<   dS )	FramerE   rz   zbytes | str | tuple[int, str]payloadrI   frame_finishedmessage_finishedNrt   r   r   r   r   r      s?         NNN****r   r   r$   r   nbytesr{   r   c                    t          |           |k    r| S | d |         } |                     dd                              d          S )Nutf-8ignore)errors)r.   decodeencode)r$   r   s     r   _truncate_utf8r      sM    
4yyF =D ;;wx;0077@@@r   c                  H    e Zd ZdddZdd	ZddZddZddZddZddZ	dS )BufferNinitial_bytesr}   r   r   c                l    t                      | _        d| _        |r|                     |           d S d S Nr   )r%   buffer
bytes_usedfeed)r"   r   s     r   r#   zBuffer.__init__  s>    kk 	%IIm$$$$$	% 	%r   	new_bytesr   c                &    | xj         |z  c_         d S r   )r   )r"   r   s     r   r   zBuffer.feed  s    y r   r   r{   r%   c                    |st                      S | j        | j        | j        |z            }| xj        t          |          z  c_        |S r   )r%   r   r   r.   )r"   r   r$   s      r   consume_at_mostzBuffer.consume_at_most  sI     	;;{4?T_v-EEF3t99$r   bytearray | Nonec                p    t          | j                  | j        z
  |k     rd S |                     |          S r   )r.   r   r   r   )r"   r   s     r   consume_exactlyzBuffer.consume_exactly  s7    t{do-664##F+++r   c                2    | j         d | j        = d| _        d S r   )r   r   rK   s    r   commitzBuffer.commit  s    K)$/)*r   c                    d| _         d S r   )r   rK   s    r   rollbackzBuffer.rollback  s    r   c                *    t          | j                  S r   )r.   r   rK   s    r   __len__zBuffer.__len__!  s    4;r   r   )r   r}   r   r   )r   r   r   r   )r   r{   r   r%   )r   r{   r   r   r   r   )r   r{   )
r5   r6   r7   r#   r   r   r   r   r   r   r   r   r   r   r     s        % % % % %! ! ! !   , , , ,   
              r   r   c                      e Zd ZddZd	dZdS )
MessageDecoderr   r   c                "    d | _         d | _        d S r   )rz   decoderrK   s    r   r#   zMessageDecoder.__init__&  s    %)26r   framer   c                   |j                                         rJ | j         1|j         t          j        u rd}t	          |          |j         | _         n,|j         t          j        urd|j         }t	          |          |j         t          j        u r t          d                      | _        |j        o|j	        }| j        |j
        }n~t          |j
        t          t          f          sJ 	 | j                            |j
        |          }n9# t          $ r,}t	          t!          |          t"          j                  d }~ww xY wt'          | j         ||j        |          }|rd | _         d | _        |S )Nzunexpected CONTINUATIONzexpected CONTINUATION, got r   )rz   rL   rE   rO   rg   rP   r   r   r   r   r   
isinstancer   r%   r   UnicodeDecodeErrorri   rV   r^   r   )r"   r   rh   finishedr$   excs         r   process_framezMessageDecoder.process_frame*  sv   <))+++++;|v222/!#&&&,DKK\!444@@@Cc"""<6;&&9099;;DL'BE,B<=DDemeY-?@@@@@T|**5=(CC% T T T!#c((K,RSSST dk4)=xHH 	 DKDLs   0 D 
E'EENr   r   r   r   r   )r5   r6   r7   r#   r   r   r   r   r   r   %  s<        7 7 7 7           r   r   c                  B    e Zd Z	 dddZddZddZddZddZddZdS )FrameDecoderNclientrI   
extensionslist[Extension] | Noner   r   c                    || _         |pg | _        t                      | _        d | _        d | _        d | _        d| _        d| _        d S r   )	r   r   r   r   headereffective_opcodemaskerpayload_requiredpayload_consumedr"   r   r   s      r   r#   zFrameDecoder.__init__N  sN     $*hh%)/3>B ! !r   r$   r   c                :    | j                             |           d S r   )r   r   r;   s     r   receive_byteszFrameDecoder.receive_bytes\  s    r   Frame | Nonec                P   | j         s|                                 sd S | j         J | j        J | j        J t	          | j                  | j        k     rd S | j         j        | j        z
  }| j        	                    |          }|s| j         j        dk    rd S | j        
                                 | xj        t	          |          z  c_        | j        | j         j        k    }| j                            |          }| j        D ][}|                    | t          |                    }t          |t                     rd}t#          ||          t%          |          }\|rkt%                      }| j        D ]P}|                    | | j         j                  }t          |t                     rd}t#          ||          |||z  }Q||z  }t+          | j        t          |          || j         j                  }	|rd | _         d | _        d | _        nt,          j        | _        |	S )Nr   error in extension)r   parse_headerr   r   r.   r   r   r|   r   r   r   r3   r   frame_inbound_payload_datar   r   rV   rg   r%   frame_inbound_completerx   r   rE   rO   )
r"   payload_remainingr   r   	extensionpayload_rh   finalresultr   s
             r   process_bufferzFrameDecoder.process_buffer_  s4   { 	4#4#4#6#6 	4{&&&{&&&$000t{d3334 K3d6KK+--.?@@ 	4;2Q664W-(DK,CC+%%g.. 	* 	*I ;;D%..QQH(K00 1*!#x000))GG 		KKE!_ $ $	"99$PPfk22 3.C%c6222%VOEuGd+U7^^Xt{WW 	8DK$(D!DKK$*$7D!r   c                L   | j                             d          }|| j                                          dS t          |d         t          z            }t          t          |d         t          z            t          |d         t          z            t          |d         t          z                      }|d         t          z  }	 t          |          }n## t          $ r d|d}t          |          w xY w|                                r|sd}t          |          t          |d         t          z            }|d         t          z  }|                     ||          }|| j                                          dS |                     |||           |r| j        rd}t          |          |s| j        sd	}t          |          |rL| j                             d
          }	|	| j                                          dS t'          |	          | _        nt+                      | _        | j                                          t/          ||||d           | _        | j        j        | _        | j        j                                        r|| _        nd| _        d| _        dS )Nr+   Fr   zInvalid opcode z#xz)Invalid attempt to fragment control framer
   z'client received unexpected masked framez)server received unexpected unmasked framer*   T)r   r   r   rI   FIN_MASKrp   	RSV1_MASK	RSV2_MASK	RSV3_MASKOPCODE_MASKrE   
ValueErrorrg   rL   	MASK_MASKPAYLOAD_LEN_MASKparse_extended_payload_lengthextension_processingr   r   r   r9   r   rw   r   rz   r   r   r   )
r"   r$   rx   ry   rz   rh   has_maskpayload_len_shortr|   r   s
             r   r   zFrameDecoder.parse_header  s   {**1--<K  """547X%&&a9$%%a9$%%a9$%%
 

 a;&	#F^^FF 	# 	# 	#/F///Cc"""	#  	#c 	#=Cc"""Q)+,, G&6688ARSSK  """5!!&#{;;; 	# 	#;Cc""" 	# 	#=Cc""" 	*+55a88K"$$&&&u)+66DKK'//DKS#v{DAA $ 2;'')) 	&$/D!!$%D! !ts   C  C8rz   rE   r|   r{   
int | Nonec                   |                                 r|t          k    rd}t          |          |t          k    rR| j                            d          }|d S t          j        d|          \  }|t          k    rd}t          |          nr|t          k    rg| j                            d          }|d S t          j        d|          \  }|t          k    rd}t          |          |dz	  rd	}t          |          |S )
Nz$Control frame with payload len > 125r+   !Hz6Payload length used 2 bytes when 1 would have sufficedrF   !Qz6Payload length used 8 bytes when 2 would have sufficed?   z'8-byte payload length with non-zero MSB)
rL   MAX_PAYLOAD_NORMALrg   PAYLOAD_LENGTH_TWO_BYTEr   r   structunpackPAYLOAD_LENGTH_EIGHT_BYTEMAX_PAYLOAD_TWO_BYTE)r"   rz   r|   rh   r$   s        r   r   z*FrameDecoder.parse_extended_payload_length  s0     	#+0B"B"B8Cc"""111;..q11D|t#]466N[000N!   1
 555;..q11D|t#]466N[222N!   b  ' @!#&&&r   ry   rp   c                .   g d}| j         D ]]}|                    | |||          }t          |t                    rd}t	          ||          t          |          D ]\  }}	|	rd||<   ^t          ||          D ]\  }
}|r|
sd}t	          |          d S )N)FFFr   TzReserved bit set unexpectedly)r   frame_inbound_headerr   rV   rg   	enumeratezip)r"   rz   ry   r|   rsv_usedr   r   rh   bitusedexpectedfounds               r   r   z!FrameDecoder.extension_processing  s     )(( 	) 	)I33D&#{SSF&+.. /*!#v...&v.. ) )	T )$(HSM)  #8S11 	' 	'OHe 'X '5!#&&&	' 	'r   r   )r   rI   r   r   r   r   r$   r   r   r   )r   r   rM   )rz   rE   r|   r{   r   r   )rz   rE   ry   rp   r|   r{   r   r   )	r5   r6   r7   r#   r   r   r   r   r   r   r   r   r   r   M  s        AE" " " " "   1 1 1 1f7 7 7 7r       D' ' ' ' ' 'r   r   c                  t    e Zd Zd)dZd*d
Zd+dZd,dZd-dZd.d/dZd0d1dZ	d0d1dZ
	 d2d3d Zd4d&Z	 d2d5d(ZdS )6FrameProtocolr   rI   r   list[Extension]r   r   c                    || _         d |D             | _        t          | j         | j                  | _        t	                      | _        |                                 | _        d | _        d S )Nc                :    g | ]}|                                 |S r   enabledr   exts     r   r   z*FrameProtocol.__init__.<locals>.<listcomp>  s%    FFF3F3FFFr   )	r   r   r   _frame_decoderr   _message_decoder_parse_more_gen_parse_more_outbound_opcoder   s      r   r#   zFrameProtocol.__init__   sf    FF*FFF +4;HH . 0 0//11/3r   r   r   c                Z   |j         }t          |t          t          f          sJ |st          j        df}nTt          |          dk    rd}t          |          t          j	        d|d d                   \  }|t          k     s|t          k    rd}t          |          t          j        t                    5  t	          |          }d d d            n# 1 swxY w Y   |t          v rd}t          |          t          |t                    s|t           k    rd}t          |          	 |dd                              d	          }n<# t$          $ r/}t          d
t'          |          z   t          j                  d }~ww xY w||f}t+          |j        ||j        |j                  S )N r
   zCLOSE with 1 byte payloadr   r+   zCLOSE with invalid codez#remote CLOSE with local-only reasonz CLOSE with unknown reserved coder   zError decoding CLOSE reason: )r   r   r   r%   rV   r\   r.   rg   r   r   MIN_CLOSE_REASONMAX_CLOSE_REASON
contextlibsuppressr   LOCAL_ONLY_CLOSE_REASONSMAX_PROTOCOL_CLOSE_REASONr   r   ri   r^   r   rz   r   r   )r"   r   r$   rh   rj   reasonr   s          r   _process_closezFrameProtocol._process_close  s   }$	 233333 	"  .3DDYY!^^-Cc"""mD$rr(33GT&&&$1A*A*A/!#&&&$Z00 ) )"4(() ) ) ) ) ) ) ) ) ) ) ) ) ) )///;!#&&&dK00 'T=V5V5V8!#&&&abb11%   !3c#hh>:  
 &>DU\4)=u?UVVVs*   ;CCC.E 
F*F  F#Generator[Frame | None, None, None]c              #  D  K   d | j         D             | _         d}|s| j                                        }|`|j                                        s| j                            |          }n,|j        t          j        k    r| 	                    |          }d}|V  |d S d S )Nc                :    g | ]}|                                 |S r   r   r   s     r   r   z1FrameProtocol._parse_more_gen.<locals>.<listcomp>4  s%    KKK3S[[]]K3KKKr   FT)
r   r   r   rz   rL   r   r   rE   rR   r   )r"   closedr   s      r   r   zFrameProtocol._parse_more_gen.  s       LK$/KKK 
	'6688E |--// " 1??FFEE\V\11 //66E!FKKK  
	 
	 
	 
	 
	r   r$   r   c                :    | j                             |           d S r   )r   r   r;   s     r   r   zFrameProtocol.receive_bytesC  s    ))$/////r   Generator[Frame, None, None]c              #  0   K   | j         D ]}| d S |V  d S r   )r   )r"   events     r   received_frameszFrameProtocol.received_framesF  s9      % 	 	E}		 	r   Nrj   r   r   
str | Noner%   c                   t                      }|t          j        u rd }||rd}t          |          |t          v rt          j        }|U|t          t          j        d|                    z  }|.|t          |	                    d          t          dz
            z  }|                     t          j        |          S )Nz&cannot specify a reason without a coder   r   r+   )r%   rV   r\   	TypeErrorr   rX   r   packr   r   r   _serialize_framerE   rR   )r"   rj   r   r   rh   s        r   closezFrameProtocol.closeM  s    ++;---D<F<:CC.. +++-DyT4!8!8999G!>MM'**,>,B   $$V\7;;;r   r   r   c                B    |                      t          j        |          S r   )r  rE   rS   r"   r   s     r   pingzFrameProtocol.ping_      $$V[':::r   c                B    |                      t          j        |          S r   )r  rE   rT   r  s     r   pongzFrameProtocol.pongb  r  r   Tbytes | bytearray | strrx   c                   t          |t          t          t          f          rt          j        }nHt          |t                    r"t          j        }|                    d          }nd}t          |          | j
        || _
        n&| j
        |urd}t          |          t          j        }|rd | _
        |                     |||          S )Nr   zMust provide bytes or textz!Data type mismatch inside message)r   r   r%   
memoryviewrE   rQ   ri   rP   r   r  r   rO   r  )r"   r   rx   rz   rh   s        r   	send_datazFrameProtocol.send_datae  s     gy*=>> 	!]FF%% 	![FnnW--GG.CC..  ($*D!!"&005CC.. (F 	)$(D!$$VWc:::r   ry   rp   rz   rE   r{   c                    t          |          dz  }t          |j                  dz  t          |j                  dz  z   t          |j                  dz  z   }t          |          }||z  |z  S )N         r*   )r{   rq   rr   rs   )r"   rx   ry   rz   fin_bitsrsv_bitsopcode_bitss          r   _make_fin_rsv_opcodez"FrameProtocol._make_fin_rsv_opcode~  sc    s88q=MMQ&3sx==A+=>#ch--STBTU&kk("[00r   bytes | bytearrayc           	        t          |          }t          ddd          }t          | j                  D ]+}|                    | ||t          |          |          \  }},|                     |||          }t          |          }d}|t          k    r|}	d }
n |t          k    r
t          }	|}
nt          }	|}
d}| j        r|	dz  }	t          ||	g          }|
r|                                rd}t          |          |r&|t          t          j        d|
                    z  }n%|t          t          j        d|
                    z  }| j        ret          t#          j        d                    }t'          |          }t          ||z   |                    t          |                    z             S t          ||z             S )NFTr?   z"payload too long for control framer   r   r*   )r%   rp   reversedr   frame_outboundr   r!  r.   r   r   r   r   r   rL   r   r   r  osurandomr   r3   )r"   rz   r   rx   ry   r   fin_rsv_opcodepayload_lengthquad_payloadfirst_payloadsecond_payloadr   rh   r   r   s                  r   r  zFrameProtocol._serialize_frame  s    G$$eUE**!$/22 	\ 	\I$33D&#uW~~WZ[[LC223VDDW///*M!NN3333M+NN5M+NL; 	$V#MNM:;;%!! &: oo% G)FKn$E$EFFF)FKn$E$EFFF; 	X $BJqMM22K$[11FVk1FNN9WCUCU4V4VVWWW')***r   )r   rI   r   r   r   r   r   )r   r  r   )r   r  )NN)rj   r   r   r
  r   r%   )r   )r   r   r   r%   )r   T)r   r  rx   rI   r   r%   )rx   rI   ry   rp   rz   rE   r   r{   )rz   rE   r   r"  rx   rI   r   r%   )r5   r6   r7   r#   r   r   r   r	  r  r  r  r  r!  r  r   r   r   r   r     s       	4 	4 	4 	4!W !W !W !WF   *0 0 0 0   < < < < <$; ; ; ; ;; ; ; ; ; CG; ; ; ; ;21 1 1 1 MQ5+ 5+ 5+ 5+ 5+ 5+ 5+r   r   )r$   r   r   r{   r   r   )=rN   
__future__r   r   r&  r   codecsr   r   enumr   typingr   r   collections.abcr	   r   r   r   r(   r   r9   r   r   r   r   MAX_PAYLOAD_EIGHT_BYTEMAX_FRAME_PAYLOADr   r   r   r   r   r   r   rE   rV   r\   r]   re   r   r   MIN_PROTOCOL_CLOSE_REASONr   MIN_LIBRARY_CLOSE_REASONMAX_LIBRARY_CLOSE_REASONMIN_PRIVATE_CLOSE_REASONMAX_PRIVATE_CLOSE_REASONr   r  	NULL_MASK	Exceptionrg   rp   rw   r   r   r   r   r   r   r   r   r   <module>r;     s    # " " " " "     				  < < < < < < < <       , , , , , , , , &))))))%%%%%% DCc

CCC
       2             " *  	  			! ! ! ! !W ! ! !6T  T  T  T  T ' T  T  T r  $             FKa  	    )       j       Z       J   A A A A!  !  !  !  !  !  !  ! H% % % % % % % %Po' o' o' o' o' o' o' o'd{+ {+ {+ {+ {+ {+ {+ {+ {+ {+r   