# to install websocket module: pip install websocket-client
import websocket
from argparse import ArgumentParser
from argparse import RawTextHelpFormatter

def websocket_downloader(uri, origin, extra_header, first_message):
    # websocket.enableTrace(True)
    ws = websocket.WebSocket()
    ws.connect(uri, origin=origin, header=extra_header)

    ws.send_binary(first_message.encode('UTF-8'))
    print('sent message: ' + first_message)

    rev = ws.recv_frame()
    rev_str = rev.data.decode('UTF-8')
    print('recieved message: ' + rev_str)

    rev_arr = rev_str.split('_')
    total_size = int(rev_arr[1])
    segment_size = int(rev_arr[2])
    print(f'total_size: {total_size}')
    print(f'segment_size: {segment_size}')

    data = bytearray()
    next_frame = 0
    while next_frame < total_size:
        print(f'{next_frame}/{total_size}',end='\r')
        ws.send_binary(str(next_frame).encode('UTF-8'))
        k = 0
        while True:
            rev = ws.recv_frame()
            rev_bytes = len(rev.data)
            if rev_bytes > 1:
                k += 1
                pos = rev.data.find('_'.encode('UTF-8'))
                pos = rev.data.find('_'.encode('UTF-8'), pos+1)
                data.extend(rev.data[pos+1:])
            else:
                break
        next_frame += k * segment_size
        
    ws.close()
          
    filename = first_message + '.mp4'
    with open(filename,'wb') as f:
        f.write(data)
    print(f'{len(data)}/{total_size}\nsaved as {filename}')


if __name__ == '__main__':
    # build ArgumentParser
    parser = ArgumentParser(description='example: %(prog)s -u "wss://simonwright.xyz/" -m "8inbm9Z8Bo_sd" -g "https://geoip.redirect-ads.com"\n\nhydrax websocket downloader v1.210516 by tamagoyaki@videohelp', formatter_class=RawTextHelpFormatter)
    # websocket uri, for exmaple: "wss://simonwright.xyz/"
    parser.add_argument('-u', '--uri', help='websocket uri, for exmaple: "wss://simonwright.xyz/"', metavar='URI', dest='ws_uri')
    # first message to be sent, for example:"8inbm9Z8Bo_sd"
    parser.add_argument('-m', '--first_message', help='first message to be sent, for example:"8inbm9Z8Bo_sd"', metavar='MESSAGE', dest='first_message')
    # origin url, for example: "https://geoip.redirect-ads.com"
    parser.add_argument('-g', '--origin', help='set origin url, for example: "https://geoip.redirect-ads.com"', metavar='ORIGIN', dest='origin_url')
    
    extra_header = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/90.0.4430.212 Safari/537.36"}

    origin = "https://geoip.redirect-ads.com"

    uri = "wss://simonwright.xyz/"

    first_message = '8inbm9Z8Bo_sd'    
    
    # handle args
    args = parser.parse_args()
    if args.ws_uri and args.first_message and args.origin_url:
        uri = args.ws_uri
        first_message = args.first_message
        origin = args.origin_url
        
        websocket_downloader(uri, origin, extra_header, first_message)

    else:
        parser.print_help()







