o
    iV                     @   sj  d dl Z d dlmZ ddlmZ ddlmZ ejr"dd Zdd	 Z	nd
d Zdd	 Z	dd Z
dZdd Zdd Zdd ZdOddZdOddZdd Zdd Zdd Zdd d!d"Zd#d$ Zd%d& Zd'd( Zd)d* Zd+Zd,Zd-Zd.Zd/Zd0Zd1Zd2d3 Z d4d5 Z!d6d7 Z"d8d9 Z#d:d; Z$d<d= Z%d>d? Z&d@dA Z'dBdC Z(dDdE Z)dFdG Z*dHdI Z+dJdK Z,dLdM Z-g dNZ.dS )P    N)ceil   
compat_ord)
Cryptodomec                 C   s   t j|t jj|| S )z/ Decrypt bytes with AES-CBC using pycryptodome )r   AESnewMODE_CBCdecryptdatakeyiv r   >/home/ubuntu/.local/lib/python3.10/site-packages/yt_dlp/aes.pyaes_cbc_decrypt_bytes      r   c                 C   s   t j|t jj|| |S )z/ Decrypt bytes with AES-GCM using pycryptodome )r   r   r   MODE_GCMdecrypt_and_verifyr   r   tagnoncer   r   r    aes_gcm_decrypt_and_verify_bytes   s   r   c                 C   s   t ttt| ||f S )zZ Decrypt bytes with AES-CBC using native implementation since pycryptodome is unavailable )bytesaes_cbc_decryptmaplistr   r   r   r   r      s   c                 C   s   t ttt| |||f S )zZ Decrypt bytes with AES-GCM using native implementation since pycryptodome is unavailable )r   aes_gcm_decrypt_and_verifyr   r   r   r   r   r   r      r   c                 K   s   t ttt| ||fi |S N)r   aes_cbc_encryptr   r   )r   r   r   kwargsr   r   r   aes_cbc_encrypt_bytes   s   r!      c                 C   s   | d t | d   S )Nr   r   r   r   r   unpad_pkcs7!      r%   c                 C   s   t t| t   }| |g|  S )zn
    PKCS#7 padding

    @param {int[]} data        cleartext
    @returns {int[]}           padding data
    )BLOCK_SIZE_BYTESlen)r   remaining_lengthr   r   r   pkcs7_padding%   s   r*   c                 C   st   t t|  }|dddd}|dk rtd||vr!td| d|dkr1|r1g | d} |d	8 }| || g|  S )
z
    Pad a block with the given padding mode
    @param {int[]} block        block to pad
    @param padding_mode         padding mode
    r       )pkcs7iso7816
whitespacezerozBlock size exceededzPadding mode z is not implementedr-      r   )r'   r(   
ValueErrorNotImplementedError)blockpadding_modepadding_sizePADDING_BYTEr   r   r   	pad_block1   s   r7   c                 C   sX   t |}tt| t }g }t|D ]}| |t |d t  }|tt||7 }q|S )a   
    Encrypt with aes in ECB mode. Using PKCS#7 padding

    @param {int[]} data        cleartext
    @param {int[]} key         16/24/32-Byte cipher key
    @param {int[]} iv          Unused for this mode
    @returns {int[]}           encrypted data
    r   )key_expansionr   r(   r'   rangeaes_encryptr*   r   r   r   expanded_keyblock_countencrypted_datair3   r   r   r   aes_ecb_encryptL   s   	r@   c                 C   s`   t |}tt| t }g }t|D ]}| |t |d t  }|t||7 }q|dt|  S )z
    Decrypt with aes in ECB mode

    @param {int[]} data        cleartext
    @param {int[]} key         16/24/32-Byte cipher key
    @param {int[]} iv          Unused for this mode
    @returns {int[]}           decrypted data
    r   N)r8   r   r(   r'   r9   aes_decryptr;   r   r   r   aes_ecb_decrypt`   s   	rB   c                 C   s   t | ||S )z
    Decrypt with aes in counter mode

    @param {int[]} data        cipher
    @param {int[]} key         16/24/32-Byte cipher key
    @param {int[]} iv          16-Byte initialization vector
    @returns {int[]}           decrypted data
    )aes_ctr_encryptr   r   r   r   aes_ctr_decrypts   s   	rD   c                 C   s   t |}tt| t }t|}g }t|D ])}t|}| |t |d t  }	|	dgtt|	  7 }	t||}
|t|	|
7 }q|dt|  S )z
    Encrypt with aes in counter mode

    @param {int[]} data        cleartext
    @param {int[]} key         16/24/32-Byte cipher key
    @param {int[]} iv          16-Byte initialization vector
    @returns {int[]}           encrypted data
    r   r   N)	r8   r   r(   r'   iter_vectorr9   nextr:   xor)r   r   r   r<   r=   counterr>   r?   counter_blockr3   cipher_counter_blockr   r   r   rC      s   	
rC   c           
      C   s   t |}tt| t }g }|}t|D ]'}| |t |d t  }|dgtt|  7 }t||}	|t|	|7 }|}q|dt|  S )z
    Decrypt with aes in CBC mode

    @param {int[]} data        cipher
    @param {int[]} key         16/24/32-Byte cipher key
    @param {int[]} iv          16-Byte IV
    @returns {int[]}           decrypted data
    r   r   N)r8   r   r(   r'   r9   rA   rG   )
r   r   r   r<   r=   decrypted_dataprevious_cipher_blockr?   r3   decrypted_blockr   r   r   r      s   	
r   r,   )r4   c                C   st   t |}tt| t }g }|}t|D ]#}| |t |d t  }	t|	|}	t|	|}
t|
|}||7 }|}q|S )a  
    Encrypt with aes in CBC mode

    @param {int[]} data        cleartext
    @param {int[]} key         16/24/32-Byte cipher key
    @param {int[]} iv          16-Byte IV
    @param padding_mode        Padding mode to use
    @returns {int[]}           encrypted data
    r   )r8   r   r(   r'   r9   r7   rG   r:   )r   r   r   r4   r<   r=   r>   rL   r?   r3   mixed_blockencrypted_blockr   r   r   r      s   



r   c              	   C   s  t dgt t|}t|dkrg |dddd}n$tt|t  t d }|dg|  tdt| dd }t||}t|}t| ||dgtt|   }	tt| t  t }
t|| dg|
  tdddt| d dd  }|t	|||krt
d|	S )a9  
    Decrypt with aes in GBM mode and checks authenticity using tag

    @param {int[]} data        cipher
    @param {int[]} key         16-Byte cipher key
    @param {int[]} tag         authentication tag
    @param {int[]} nonce       IV (recommended 12-Byte)
    @returns {int[]}           decrypted data
    r      r      bigzMismatching authentication tag)r:   r'   r8   r(   r   to_bytesghashincrD   rC   r1   )r   r   r   r   hash_subkeyj0fillghash_iniv_ctrrK   pad_lens_tagr   r   r   r      s,   &
r   c                 C   s~   t |t d }t| |dt } td|d D ]$}t| } t| } ||kr-tt| t} t| ||t |d t  } q| S )z
    Encrypt one block with aes

    @param {int[]} data          16-Byte state
    @param {int[]} expanded_key  176/208/240-Byte expanded key
    @returns {int[]}             16-Byte cipher
    r   N)	r(   r'   rG   r9   	sub_bytes
shift_rowsr   iter_mix_columnsMIX_COLUMN_MATRIXr   r<   roundsr?   r   r   r   r:      s    r:   c                 C   sx   t |t d }t|ddD ]$}t| ||t |d t  } ||kr*tt| t} t| } t| } qt| |dt S )z
    Decrypt one block with aes

    @param {int[]} data          16-Byte cipher
    @param {int[]} expanded_key  176/208/240-Byte expanded key
    @returns {int[]}             16-Byte state
    r   r   r#   N)	r(   r'   r9   rG   r   r_   MIX_COLUMN_MATRIX_INVshift_rows_invsub_bytes_invra   r   r   r   rA     s   
rA   c                 C   s   d}t t| } t | }|d| dg|t|   }t|dt t||t  }| d| }| |d }t|||dgt|   }t	|S )af  
    Decrypt text
    - The first 8 Bytes of decoded 'data' are the 8 high Bytes of the counter
    - The cipher key is retrieved by encrypting the first 16 Byte of 'password'
      with the first 'key_size_bytes' Bytes from 'password' (if necessary filled with 0's)
    - Mode of operation is 'counter'

    @param {str} data                    Base64 encoded string
    @param {str,unicode} password        Password (will be encoded with utf-8)
    @param {int} key_size_bytes          Possible values: 16 for 128-Bit, 24 for 192-Bit or 32 for 256-Bit
    @returns {str}                       Decrypted data
    rQ   Nr   )
r   base64	b64decodeencoder(   r:   r'   r8   rD   r   )r   passwordkey_size_bytesNONCE_LENGTH_BYTESr   r   cipherrK   r   r   r   aes_decrypt_text  s   rm   )   r         rQ   r"   r+   @   r0      6   (   c   |   w   {      k   o      0   r   g   +            v            }      Y   G                        r               &   rs   ?         4            q      1      rp      #                        r0         '      u   	      ,      rr   n   Z      R   ;         )      /      S      r      r+         [   j         9   J   L   X                  C   M   3      E      ro      P   <         Q      rq            8               !   r"               rP         _      D            ~   =   d   ]      s   `      O      "   *         F               ^            2   :   
   I      $   \            b            y         7   m   rn      N      l   V         e   z      rQ      x   %   .                     t      K            p   >      f   H            a   5   W                              i                           U   (                        B   h   A      -         T         (   r   r   r   r  r|   rs   r   r   r[  rq   r   rF  r   r   r   r   ru   r   r   r   rO  r   r   rQ  r   rM  r   r   r   r  rR  r   rd  rw   rN  r
  r,  r  r   r   r  r   r  r  r]  r   r   r  rQ   r*  rX  r:  rU  rL  r  r   r   r   r   r  r  r5  r   r)  r   rH  r=  r   rC  r^  rI  rf  r   r   r  r   r   r$  r   r   r   r7  r;  r   r   r   rB  r   r  r   r  rA  r   rn   r   r   r   r   r   r   rW  r   r  r  r   r  r   r   r  r   r   r  r   r   rP  r   r   r   rb  ro   rD  r   r4  r<  r   r   r6  ry   r  r  rJ  r_  r   r}   r   r#  r   rx   r   rS  r   r-  r\  r   r   r  r1  r   r  r   r@  r   r   r   r  r/  r+  r   rV  r   r   r   r   r   rE  r   r{   rY  rz   r   r  r>  r   r   r   rr   r   r!  r8  r3  r.  r   r  r+   r   r  r   r   r(  r   r   r"  r2  r0  r   r   r   r   r   r   r   r   r"   r   r   r0   r   r   r   r   r   r  r   r9  r   rZ  ra  r   r%  r   r   r   r   r   r   r	  r   r   r&  r   r   rc  r  r   re  r   r   r   r`  r?  r   r~   rp   r   r'  rv   r   r   rG  rK  r  rt   rT  r   rP   r   ))ro   r<  r   r   )r   ro   r<  r   )r   r   ro   r<  )r<  r   r   ro   ))r>  r  rZ  r   )r   r>  r  rZ  )rZ  r   r>  r  )r  rZ  r   r>  (   r   r<  r   rb  rJ  r   rT  r   r   r*  r   r   rX  rH  r   r@  r   rG  r   r;  r   r   r  r   r   ro   r  r  rP  r   r:  r   r   r   r  r  r  r   r   r   r   r   rL  r7  r   r   r\  r   r   r   rp   rP   r  r   r   r   r   r   r^  r  r  r   r   r   r   r   r}   r  r	  r   r   r   r  r,  r   rQ   r   rU  r(  r   r   rF  rB  r   ry   r4  r   r   r   rI  r   rS  r  r  r   r   r9  r   rA  r   r"   r|   r   r   r  rE  r   rK  re  r   r?  r   r   r   r~   r   rQ  r   r   r   r   r   r   r&  rR  r+   r   r   r   rf  r  r  r   r  r   r  r   r  r
  r!  r   r   r   r_  r   r  r   r   r   r   rq   r   r   r   r   r1  r   r[  r   r   r   r'  r  r   r  r   r   r   r   r   r   rV  r%  rM  rY  r0   rO  r   rD  r   r/  r   r$  r   r#  r)  rz   r   r  r   r{   rd  r   r2  r   rt   r   r"  r   r   rr   ra  rv   r`  rc  r   r  r   r   r   r   r  r  r5  rC  r  r   r   r8  r]  r.  r   r   r>  r   rs   r   r  r   rw   rn   rW  r   r6  r   rN  r   rx   rZ  r   r   r3  r0  ru   r   r   r   r   r+  r  r   r-  r   r   r=  r   (   r   r   r   r   r
  ro   r   r.  r3  r   rr   r^  r   r  rV  r<  r   rp   r	  r>  r   rn   r   r   r   r   rQ   r  rH  rK  r+  rD  r   r  rE  r9  r   rB  r   r   r   r  r,  r   r   r   r   r(  r$  r   r6  r   r   rb  rG  r  r   r   r   r   r@  r   r   rM  r   r   r  r4  rs   r   rS  rN  r   r  r   r   rq   r  r   r   r:  r0  r   r|   r[  r  r5  r  r   r)  r   rI  r   r   r  r"   r   r   r;  r   r   r   rP  r]  r  ry   rU  rd  r   r   r   r'  r~   r  r  r   rO  r   r  r   r  r   r  r   r   r   r   rA  r   r   r   r   r"  r#  r   r1  r   r&  rR  r  r  r\  r   r/  r   r   r   r%  r   rf  r  r   r   r   r   rc  r   r  r   r   r   rP   r=  rz   r   r   r  r   r   r   r2  ra  r   r   rw   r   r   re  r8  r   r   r   r   rC  r   r   rX  r   r   rT  r   r   r   r   rQ  r   r?  r   r   r   r   r  r   r   r  r   r   r   r   r   r   r   r_  r   r  r   r  r   rF  r   r!  rx   r  r   r   rJ  r   rL  r   r+   r*  rY  r-  ru   r  r   rv   r`  r   r   r}   r   r   r  r{   r   r   r   rZ  rt   rW  r0   r   r   r7  r   c                 C   sB  | dd } d}t | }|d d t }t | |k r| dd }t||}|d7 }| t|| | d|  7 } tdD ]}| dd }| t|| | d|  7 } q;|dkrn| dd }t|}| t|| | d|  7 } t|dkrudn|dkr{d	nd
D ]}| dd }| t|| | d|  7 } q~t | |k s| d| S )z
    Generate key schedule

    @param {int[]} data  16/24/32-Byte cipher key
    @returns {int[]}     176/208/240-Byte expanded key
    Nr   rp   r   r<  r+   r   ro   r   )r(   r'   key_schedule_corerG   r9   r]   )r   rcon_iterationrj   expanded_key_size_bytestemp_r   r   r   r8     s*   
$r8   c                 c   s    	 | V  t | } qr   )rU   )r   r   r   r   rE     s
   rE   c                 C      dd | D S )Nc                 S      g | ]}t | qS r   )SBOX.0xr   r   r   
<listcomp>      zsub_bytes.<locals>.<listcomp>r   r$   r   r   r   r]        r]   c                 C   rm  )Nc                 S   rn  r   )SBOX_INVrp  r   r   r   rs    rt  z!sub_bytes_inv.<locals>.<listcomp>r   r$   r   r   r   re     ru  re   c                 C   s   g | dd  | d S )Nr   r   r   r$   r   r   r   rotate     rw  c                 C   s(   t | } t| } | d t| A | d< | S )Nr   )rw  r]   RCON)r   ri  r   r   r   rh    s   rh  c                 C   s   dd t | |ddD S )Nc                 S   s   g | ]\}}||A qS r   r   )rq  rr  yr   r   r   rs    s    zxor.<locals>.<listcomp>F)strict)zip)data1data2r   r   r   rG     rx  rG   c              	   c   s    dD ];}|D ]6}d}t dD ]*}|| ||d  | dks$|| dkr&dntt| ||   t||   d  N }q|V  qqd S )N)r   rp   rQ   rP   r   rp   r   )r9   RIJNDAEL_EXP_TABLERIJNDAEL_LOG_TABLE)r   matrixr?   rowmixedjr   r   r   r_     s   *"r_   c                        fddt dD S )Nc                    s2   g | ]}t d D ]} || d@ d  |  qqS rp   r<  r9   rq  columnr  r$   r   r   rs       2 zshift_rows.<locals>.<listcomp>rp   r  r$   r   r$   r   r^     r&   r^   c                    r  )Nc                    s2   g | ]}t d D ]} || d@ d  |  qqS r  r  r  r$   r   r   rs    r  z"shift_rows_inv.<locals>.<listcomp>rp   r  r$   r   r$   r   rd     r&   rd   c                 C   s<   g }d}| D ]}|r|dO }|d@ }|dL }| | q|S )Nr      r   )append)r   data_shiftedbitnr   r   r   shift_block  s   r  c                 C   sT   | d d  } t t| d ddD ]}| | dkrd| |< q| | d | |<  | S | S )Nr   r#   r   r   )r9   r(   )r   r?   r   r   r   rU     s   
rU   c                 C   s   t | tkst |tkrtdt ddgdgtd   }|d d  }dgt }| D ]'}tdddD ]}|d|> @ r@t||}|d d@ }t|}|rQt||}q3q+|S )NzLength of blocks need to be  bytesrG  r   r   r   r#   )r(   r'   r1   r9   rG   r  )block_xblock_yblock_rblock_vblock_zr?   r  do_xorr   r   r   block_product  s    


	r  c                 C   s`   t |t rtdt ddgt }tdt |tD ]}|||t  }tt||| }q|S )NzLength of data should be r  r   )r(   r'   r1   r9   r  rG   )subkeyr   last_yr?   r3   r   r   r   rT     s   
rT   )r   r   r   r!   rD   rC   rA   rm   rB   r@   r:   r   r   r8   r7   r*   r%   r   )/rf   mathr   compatr   dependenciesr   r   r   r   r!   r'   r%   r*   r7   r@   rB   rD   rC   r   r   r   r:   rA   rm   ry  ro  rv  r`   rc   r  r  r8   rE   r]   re   rw  rh  rG   r_   r^   rd   r  rU   r  rT   __all__r   r   r   r   <module>   s\    


+!