pwdlBackendv2 / app.py
AkshitShubham's picture
Upload folder using huggingface_hub
dbe0ac5 verified
import re
import subprocess
from flask import Flask, request, jsonify, send_file
from flask_socketio import SocketIO, emit
from flask_cors import CORS
from logger import Logger
from uuid import uuid4
import os
app = Flask(__name__)
logger = Logger('log.txt')
app.config['SECRET_KEY'] = 'secret!'
CORS(app, resources={r"/*": {"origins": "*"}})
socketio = SocketIO(app, cors_allowed_origins="*")
# clients is a 'list' of all connected clients
clients = []
SIMULATE = False
TEST = False
"""
client_downloads is a 'dict'
with format as shown below
{
'clientId':
{
'downloadId':
[
[
['fileName','id'],
'done?'
]
]
}
downloadId is a unique id for each download (ref: downloadId in download.py)
downloadId is generated by the server as opposed to clientId which is generated by the client
in earlier versions of the server, downloadId was generated by the client
- these versions used static site, as opposed to the current version which uses 'react'
downloadId is further a list with a list of 2 elements
- first element is the list of fileName and id
- second element is a integer value
0 -> download has been completed
1 -> download is in progress
-1 -> download has been cancelled
"""
client_downloads = {}
def shell(command, filter='.*', clientId='None', stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
universal_newlines=True):
import os, time
# Set PYTHONUNBUFFERED environment variable
os.environ['PYTHONUNBUFFERED'] = '1'
command = to_list(command)
count = 0
process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True,
shell=True)
output_list = []
# Read and print the output in real-time
while True:
output = process.stdout.readline()
if output == '' and process.poll() is not None:
break
if output and filter is not None and re.search(filter, output):
print(output.strip())
if output not in output_list:
socketio.emit(f'progress-{clientId}', {
'message': output.strip(),
'count': count
})
output_list.append(output.strip())
count += 1
# Wait for the process to complete and get the return code
return_code = process.poll()
return return_code
def to_list(variable):
if isinstance(variable, list):
return variable
elif variable is None:
return []
else:
# Convert to string and then to list by splitting at whitespaces
return variable.split()
def check_for_existing_running_downloads(clientId, downloadId=None):
# if dowbloadId is not provided, check if clientId exists in client_downloads
# and then check for any downloadId and if has a download
# Case 0: clientId is not in clients
if clientId not in clients:
return False
# Case 1: clientId is not in client_downloads
if clientId not in client_downloads:
return False
# Case 2: [Ideal] clientId provided and downloadId is provided
if downloadId:
downloads_of_client = client_downloads[clientId]
# Case 2.1: downloadId is not in client_downloads
if downloadId not in downloads_of_client:
return False
downloads_by_downloadId = downloads_of_client[downloadId]
for download in downloads_by_downloadId:
if download[1] == 1:
return downloadId
# Case 3: clientId provided but downloadId is not provided
for downloadId in client_downloads[clientId]:
downloads_by_downloadId = client_downloads[clientId][downloadId]
for download in downloads_by_downloadId:
if download[1] == 1:
return downloadId
# Case 4 : No downloads are running
return False
def send_existing_download_status(clientId):
emit(f'exists-running-downloads-{clientId}', {
'message': 'Client has existing downloads',
'names': [download[0][0] for download in
client_downloads[clientId][check_for_existing_running_downloads(clientId)]],
'ids': [download[0][1] for download in
client_downloads[clientId][check_for_existing_running_downloads(clientId)]],
})
@socketio.on('connect')
def connect(data):
pass
@socketio.on('connect-with-client-Id')
def connect_with_client_id(data):
clientId = data['clientId']
if not clientId:
logger.logError(f'Empty CLient ID by {request.sid}')
if clientId not in clients:
logger.logInfo(f'Client {clientId} connected')
clients.append(clientId)
if not clientId in client_downloads:
client_downloads[clientId] = {
}
if check_for_existing_running_downloads(clientId):
logger.logInfo(f'Client {clientId} has existing downloads')
send_existing_download_status(clientId)
return # return as the client has an existing download
else:
logger.logInfo(f'Client {clientId} has no existing downloads')
emit(
f'successful-{clientId}',
{
'connected': True
}
)
@socketio.on('test-pwdl')
def test_pwdl(data):
names = data['names']
ids = data['IDs']
clientId = data['clientId'] # Assuming you're passing the clientId from the client
existing_download_id = check_for_existing_running_downloads(clientId) # Check if there are any existing downloads
if existing_download_id:
download_id = existing_download_id
# normally this should not happen as the client should not be able to request a download
# and the existing download is checked when the client is connected
# so this is more of a fail-safe
# this means that the client has an existing download
# hence we emit the message to the client
send_existing_download_status(clientId)
return # return as the client has an existing download
else:
download_id = str(uuid4())
client_downloads[clientId] = {
download_id: []
}
# temporarily set the download id as '13cac2f0-e46f-465b-8a2c-71a043b05bd0'
if TEST:
download_id = '13cac2f0-e46f-465b-8a2c-71a043b05bd0'
logger.logInfo(f'Client {clientId} requested download')
logger.logInfo(f'Client {clientId} requested download with assigned downloadId {download_id}')
# the download path is ./webdl/<clientId>/<downloadId>
# make the directory
try:
os.makedirs(f'./webdl/{clientId}/{download_id}', exist_ok=True)
except Exception as e:
logger.logError(f'Error creating directory {e} for client {clientId} with downloadId {download_id} ')
print(e)
return
for name, id in zip(names, ids):
if TEST: break
# creates a new download in the client_downloads dict
client_downloads[clientId][download_id].append([[name, id], 1])
# currently just simulating download (no actual download)
# command is the pwdl command that downloads the video with the given id and name
command = ['python',
'pwdlv3/pwdl.py',
'--id', id,
'--name', name,
'--dir', f'./webdl/{clientId}/{download_id}',
'--verbose']
if SIMULATE: command.append('--simulate')
# if simulate is turned on then create a simple txt file with format 'name.txt' in the download directory
if SIMULATE:
with open(f'./webdl/{clientId}/{download_id}/{name}.txt', 'w') as f:
f.write(f'This is a simulated file for {name}')
# start the download
socketio.emit(f'started-download-{clientId}', {
"name": name,
"id": id,
"downloadId": download_id
})
# execute the command (the shell command from
# 'https://github.com/shubhamakshit/EssentialPythonFunctions/blob/main/process.py' was edited for custom usage)
return_code_after_download = shell(command,
filter='.*',
clientId=clientId,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
universal_newlines=True)
print(f"Return code after download: {return_code_after_download}")
if return_code_after_download == 0:
# download has been completed
client_downloads[clientId][download_id][-1][1] = 0
socketio.emit(f'partial-complete-download-{clientId}', {
"name": name,
"id": id,
"downloadId": download_id
})
else:
# download has been cancelled
client_downloads[clientId][download_id][-1][1] = -1
socketio.emit(f'partial-cancel-download-{clientId}', {
"name": name,
"id": id,
"download_id": download_id
})
# log that downloads are complete with 'x' done and 'y' cancelled
done = []
cancelled = []
if not TEST:
for download in client_downloads[clientId][download_id]:
if download[1] == 0:
done.append(download[0][0] + (".txt" if SIMULATE else ".mp4") )
elif download[1] == -1:
cancelled.append(download[0][0])
logger.logInfo(f'Client {clientId} has completed {len(done) + 1} downloads and cancelled {len(cancelled) + 1} downloads')
# emit that the download has been completed (those not completed are cancelled)
if TEST:
done = ["1.mp4", "2.mp4"]
print(f'Download ID: {download_id}')
socketio.emit(f'download-complete-{clientId}', {
'message': 'pwdl executed',
'downloadId': download_id,
'done': done,
'cancelled': cancelled
})
@app.route('/get')
def get():
downloadId = request.args.get('downloadId')
clientId = request.args.get('clientId')
fileName = request.args.get('fileName')
print(f'DownloadId {downloadId}')
file_path = f'./webdl/{clientId}/{downloadId}/{fileName}'
response = send_file(file_path, as_attachment=True)
response.headers['Content-Length'] = os.path.getsize(file_path)
return response
@app.route('/')
def hello_world(): # put application's code here
return 'Hello World!'
if __name__ == '__main__':
socketio.run(app, debug=True, port=5001)