Untitled

from websockets.exceptions import ConnectionClosedError
from websockets import serve
import asyncio
import requests
from os.path import exists

LOG_PATH = "C:\\Users\\phant_m\\Desktop\\Chrome Extensions\\discord tracker\\log.txt"

connected_clients = set()

EXTENSIONS = [
    # Image file extensions
    ".png", ".jpg", ".jpeg", ".tiff", ".pdf", ".svg", ".gif", ".bmp", ".webp",
    ".psd", ".ai", ".eps", ".ico", ".heic", ".tif", ".raw", ".xcf", ".indd",
    ".jp2", ".jxr", ".wdp", ".hdp", ".dds", ".yuv", ".arw", ".cr2", ".dng",
    ".nef", ".rw2", ".orf", ".sr2", ".raf", ".svgz", ".erf", ".kdc", ".mrw",
    ".pef", ".dcr", ".dxf", ".3fr", ".mef", ".mos", ".nrw", ".x3f", ".crw",
    ".srf", ".srw", ".ptx", ".rwl", ".r3d", ".raw", ".rwz", ".rw1", ".rwz",
    
    # Video file extensions
    ".mp4", ".mov", ".avi", ".mkv", ".wmv", ".flv", ".webm", ".mpeg", ".mpg",
    ".m4v", ".3gp", ".3g2", ".rm", ".rmvb", ".asf", ".vob", ".ts", ".m2ts",
    ".mxf", ".ogv", ".divx", ".xvid", ".rm", ".rmvb", ".mts", ".f4v", ".dat",
    ".m2t", ".avchd", ".qt", ".swf", ".yuv", ".nsv", ".h264", ".h265", ".mod",
    ".tod", ".str", ".amv", ".m2v", ".avp", ".pva", ".pss", ".vp6", ".vp7",
    ".vp8", ".vp9", ".vro", ".mts", ".m2p", ".m1v", ".mvc", ".dav", ".mpg2",
    ".mpe", ".m2v", ".m2p", ".m2t", ".m2ts", ".mp2v", ".tpr", ".roq", ".svi",
]


def getExtension(url):
    for e in EXTENSIONS:
        if e in url:
            return e

def downloadFile(url, save_path):
    EXTENSION = getExtension(url)

    number = 0
    while exists(save_path + f'{number}{EXTENSION}'):
        number += 1

    save_path += f'{number}{EXTENSION}'
    
    # Streaming, so we can iterate over the response.
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(save_path, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)

async def updateMessages(current_socket = None):
    for socket in connected_clients:
        if socket != current_socket:
            await socket.send("messages_update:" + open(LOG_PATH, 'r', encoding='utf-8').read())

async def handler(socket):
    global connected_clients

    try:
        connected_clients.add(socket)

        text = open(LOG_PATH, 'r', encoding='utf-8').read()
        await socket.send("messages:" + text)
        await socket.send('messages_update:' + text)

        async for message in socket:
            print(message)
            try:
                if message == 'clear':
                    open(LOG_PATH, 'w').close()
                    await updateMessages()
                elif message.startswith('download:'):
                    downloadFile(message.replace('download:', ''), 'C:\\Users\\phant_m\\Desktop\\Chrome Extensions\\discord tracker\\media\\download')
                elif message != '':
                    open(LOG_PATH, 'a', encoding='utf-8').write(''.join(c for c in message if c.isprintable()) + '\n')
                    await updateMessages(socket)
            except:
                pass
    except ConnectionClosedError:
        pass
    finally:
        connected_clients.remove(socket)



async def main():
    async with serve(handler, "", 8001):
        await asyncio.Future()

if __name__ == "__main__":
    asyncio.run(main())