
    @OOf'                        d dl Z d dlmZ d dlmZ d dlZd dlZd dlZd dlm	Z	 d dlm
Z
 d dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ e G d de             Ze G d de             Ze G d de             Z G d deej<                           Z G d dej@                        Z! G d dejD                        Z# G d d      Z$y)    N)Iterator)	dataclass)ConnectionState)Opcode)
connection)http)	websocket)commands)events)layer)	StartHook)Context)MessageInjected)expectc                   0    e Zd ZU dZej
                  ed<   y)WebsocketStartHookz/
    A WebSocket connection has commenced.
    flowN__name__
__module____qualname____doc__r   HTTPFlow__annotations__     a/var/www/premiumrankchecker/venv/lib/python3.12/site-packages/mitmproxy/proxy/layers/websocket.pyr   r      s     --r   r   c                   0    e Zd ZU dZej
                  ed<   y)WebsocketMessageHooka  
    Called when a WebSocket message is received from the client or
    server. The most recent message will be flow.messages[-1]. The
    message is user-modifiable. Currently there are two types of
    messages, corresponding to the BINARY and TEXT frame types.
    r   Nr   r   r   r   r   r       s     --r   r   c                   0    e Zd ZU dZej
                  ed<   y)WebsocketEndHookzt
    A WebSocket connection has ended.
    You can check `flow.websocket.close_code` to determine why it ended.
    r   Nr   r   r   r   r!   r!   ,   s    
 --r   r!   c                       e Zd ZdZy)WebSocketMessageInjectedz;
    The user has injected a custom WebSocket message.
    N)r   r   r   r   r   r   r   r#   r#   6   s    r   r#   c                        e Zd ZU dZej
                  ed<   ee   ed<   dej
                  f fdZ	de
j                  j                  dej                  fdZd Z xZS )	WebsocketConnectionz
    A very thin wrapper around wsproto.Connection:

     - we keep the underlying connection as an attribute for easy access.
     - we add a framebuffer for incomplete messages
     - we wrap .send() so that we can directly yield it.
    conn	frame_bufc                B    t        |   |i | || _        dg| _        y )Nr   )super__init__r&   r'   )selfr&   argskwargs	__class__s       r   r*   zWebsocketConnection.__init__H   s%    $)&)	r   eventreturnc                 d    | j                  |      }t        j                  | j                  |      S N)sendr
   SendDatar&   )r+   r/   datas      r   send2zWebsocketConnection.send2M   s'    yy  D11r   c                 P    d| j                   j                   d| j                   dS )NzWebsocketConnection<z, >)statenamer&   )r+   s    r   __repr__zWebsocketConnection.__repr__Q   s$    %djjoo%6b1EEr   )r   r   r   r   r   
Connectionr   listbytesr*   wsprotor   Eventr
   r4   r6   r;   __classcell__r.   s   @r   r%   r%   <   s[     

E{J$9$9 
27>>// 2H4E4E 2Fr   r%   c                       e Zd ZU dZej
                  ed<   eed<   eed<   dedej
                  f fdZ	 e
ej                        dej                  d   fd	       ZeZ e
ej"                  ej$                  e      d
ej(                  dej                  d   fd       Z e
ej"                  ej$                  e      dej                  d   fd       Z xZS )WebsocketLayerz>
    WebSocket layer that intercepts and relays messages.
    r   	client_ws	server_wscontextc                 2    t         |   |       || _        y r2   )r)   r*   r   )r+   rG   r   r.   s      r   r*   zWebsocketLayer.__init__^   s    !	r   r0   Nc              #   *  K   g }g }| j                   j                  sJ | j                   j                  j                  j                  dd      }|rt        j
                  j                  |j                  dd            D ]  }|j                  dd      d   j                         }|t        j                  j                  j                  k(  rt        j                  j                         }|j                  |       |j                  |       t        j                  j                         }|j                  |       |j                  |       t        j                   d|d	        t#        t        j$                  j&                  || j(                  j*                  
      | _        t#        t        j$                  j.                  || j(                  j0                  
      | _        t5        | j                          | j6                  | _        y w)NzSec-WebSocket-Extensions asciireplace;   r   z%Ignoring unknown WebSocket extension .)r&   )r   responseheadersgetr?   	utilitiessplit_comma_headerencodesplitstrip
extensionsPerMessageDeflater:   finalizeappendr
   Logr%   ConnectionTypeSERVERrG   clientrE   CLIENTserverrF   r   relay_messages_handle_event)	r+   _client_extensionsserver_extensions
ext_headerextext_nameclient_deflateserver_deflates	            r   startzWebsocketLayer.startb   s     yy!!!!YY''//334NPRS
((;;!!'95  99S!,Q/557w11CCHHH%,%7%7%I%I%KN"++C0%,,^<%,%7%7%I%I%KN"++C0%,,^<",,?|1M   -""))+<4<<CVCV
 -""))+<4<<CVCV
 !++!00s   HHr/   c              #   V  K   | j                   j                  sJ t        |t        j                        r&|j
                  | j                  j                  k(  }d}n7t        |t              r|j                  j                  }d}nt        d|       |rdnd}|r| j                  }| j                  }n| j                  }| j                  }t        |t        j                        r|j                  |j                          nt        |t        j"                        r|j                  d        nt        |t              rgt%        g |j                  j&                  t(        j*                  k(        }|j,                  j/                   ||j                  j0                               nt        d|       |j                         D ]  }t        |t2        j                  j4                        rt        |j                   t6              }	|	r@t(        j*                  }
|j8                  dxx   |j                   j;                         z  cc<   n1t(        j<                  }
|j8                  dxx   |j                   z  cc<   |j>                  rdjA                  |j8                        }t%        |j8                  |	      }dg|_        t        jB                  |
|||      }| j                   j                  jD                  jG                  |       tI        | j                          |jJ                  rh ||j0                        D ]  }|jM                  |        |jN                  s|j8                  jG                  d       t        |t2        j                  jP                  t2        j                  jR                  f      rktU        jV                  d	|jX                  jZ                  j]                          d
| dt_        |j`                        d       |jM                  |       ht        |t2        j                  jb                        rJte        jd                         | j                   j                  _3        || j                   j                  _4        |jj                  | j                   j                  _6        |jn                  | j                   j                  _8        | j                  | j                  fD ]b  }|jr                  tt        jv                  tt        jx                  hv r|jM                  |       tU        jb                  |jz                         d t}        | j                          d| j                   _?        | j                  | _A        t        d|        y w)NFTzUnexpected event: r_   ra   r   )injectedzReceived WebSocket z from z (payload: )zUnexpected WebSocket event: )Br   r	   
isinstancer   ConnectionEventr   rG   r_   r#   messagefrom_clientAssertionErrorrE   rF   DataReceivedreceive_datar5   ConnectionClosedFragmentizertyper   TEXT_eventsextendcontentr?   Messagestrr'   rU   BINARYmessage_finishedjoinWebSocketMessagemessagesr[   r   droppedr6   frame_finishedPingPongr
   r\   r.   r   lowerr>   payloadCloseConnectiontimetimestamp_endclosed_by_clientcode
close_codereasonclose_reasonr9   r   OPENREMOTE_CLOSINGr&   r!   livedonerc   )r+   r/   rt   ro   from_strsrc_wsdst_wsfragmentizerws_eventis_texttypr~   rs   msgwss                  r   rb   zWebsocketLayer.relay_messages   s6    yy""""eV334**dll.A.AAKH78--33KH #5eW!=>>*8^^F^^F^^F^^FeV001

+v667%78'EMM,>,>&++,MNLNN!!,u}}/D/D"EF #5eW!=>> 5	PH(GNN$:$:;$X]]C8 ++C$$R(HMM,@,@,BB( --C$$R(HMM9(,,!hhv'7'78G#/0@0@'#JL(+uF$'88['HG II''0077@.tyy99"??#/#@ 4C"(,,s"334 ,,$$++C0Hw~~':':GNN<O<O&PQll)(*<*<*E*E*K*K*M)NfU]T^ _!!&x'7'7!8 ;1>  ll8,,Hgnn&D&DE48IIK		##17B		##419		##.3;??		##0>>4>>: <Bxx',,'66$ 
 !hhx00"22277;;< 'tyy11!&		%)YY"$'CH:%NOOk5	Ps   L&V)*8V)$IV)c              #   $   K   dE d {    y 7 w)Nr   r   )r+   rd   s     r   r   zWebsocketLayer.done   s     s   )r   r   r   r   r   r   r   r%   r   r*   r   r   Startr   CommandGeneratorrl   rc   rv   rx   r#   r@   rb   r   rA   rB   s   @r   rD   rD   U   s     --"""" t}}  FLL!1%006 !1 !1F MF!8!8:RSSPFLL SPU5K5KD5Q SP TSPj F!8!8:RS//5  Tr   rD   c                   x    e Zd ZdZdZdee   defdZdedefdZ	d	ed
e
ej                  j                     fdZy)ry   aH  
    Theory (RFC 6455):
       Unless specified otherwise by an extension, frames have no semantic
       meaning.  An intermediary might coalesce and/or split frames, [...]

    Practice:
        Some WebSocket servers reject large payload sizes.
        Other WebSocket servers reject CONTINUATION frames.

    As a workaround, we either retain the original chunking or, if the payload has been modified, use ~4kB chunks.
    If one deals with web servers that do not support CONTINUATION frames, addons need to monkeypatch FRAGMENT_SIZE
    if they need to modify the message.
    i  	fragmentsr   c                 X    |D cg c]  }t        |       c}| _        || _        y c c}w r2   )lenfragment_lengthsr   )r+   r   r   xs       r   r*   zFragmentizer.__init__   s%    1: ;AQ ; !<s   'r5   r   c                     | j                   r3|j                  d      }t        j                  j	                  ||      S t        j                  j                  ||      S )NrL   )errors)r   )r   decoder?   r   TextMessageBytesMessage)r+   r5   r   data_strs       r   r   zFragmentizer.msg   sX    <<{{){4H>>--+; .   >>..tFV.WWr   r~   r0   c              #     K   t        |      t        | j                        k(  rMd}| j                  d d D ]!  }| j                  ||||z    d       ||z  }# | j                  ||d  d       y d}t        |      | j                  z
  }||k  r9| j                  |||| j                  z    d       || j                  z  }||k  r9| j                  ||d  d       y w)Nr   rn   FT)r   sumr   r   FRAGMENT_SIZE)r+   r~   offsetfltotals        r   __call__zFragmentizer.__call__  s     w<3t4455F++CR0 hhwv<eDD" ((767+T22FL4#5#55E5.hhwv9K9K0KLeTT$,,, 5. ((767+T22s   CC!	C!N)r   r   r   r   r   r=   r>   boolr*   r   r   r?   r   r   r   r   r   r   ry   ry      s]     M$u+  X X X3 3(7>>3I3I*J 3r   ry   )%r   collections.abcr   dataclassesr   wsproto.extensionsr?   wsproto.frame_protocolwsproto.utilitiesr   r   	mitmproxyr   r   r	   mitmproxy.proxyr
   r   r   mitmproxy.proxy.commandsr   mitmproxy.proxy.contextr   mitmproxy.proxy.eventsr   mitmproxy.proxy.utilsr   r   r   r!   r   r#   r<   r%   LayerrD   ry   r   r   r   <module>r      s     $ !    # )     $ " ! . + 2 (    9   y  y/I/IJ F',, F2KU[[ K\-3 -3r   