diff --git a/deploy/setup.py b/deploy/setup.py index 10afcbb..8e71b99 100644 --- a/deploy/setup.py +++ b/deploy/setup.py @@ -138,7 +138,8 @@ def fileUpload(fileName: str) -> str: 'max-connection-per-server': '8', 'max-overall-upload-limit': '1K', 'min-split-size': '10M', 'seed-time': '0.01', 'split': '10'}, 'authorizedChats': {}, 'dlRootDir': 'dl', 'logLevel': 'INFO', 'megaAuth': {'apiKey': '', 'emailId': '', 'passPhrase': ''}, - 'statusUpdateInterval': '5', 'trackersListUrl': 'https://trackerslist.com/all_aria2.txt'} + 'statusUpdateInterval': '5', 'trackersListUrl': 'https://trackerslist.com/all_aria2.txt', + 'ytdlFormat': 'best/bestvideo+bestaudio'} envVars: typing.Dict = {'dlWaitTime': '5'} if __name__ == '__main__': diff --git a/tgmb/__init__.py b/tgmb/__init__.py index 76b76ab..720aad0 100644 --- a/tgmb/__init__.py +++ b/tgmb/__init__.py @@ -1,7 +1,5 @@ # TODO: add sufficient documentation to the functions and classes in this module -# TODO: Code for Download from Mega # TODO: Code for Upload to Mega -# TODO: Helper functions - bot_utils.py, fs_utils.py, message_utils.py # TODO: Code for user filters # TODO: Add and Handle Exceptions # TODO: Code for direct link generation @@ -210,12 +208,12 @@ def initHelper(self) -> None: self.reqVars: [str] = ['botToken', 'botOwnerId', 'telegramApiId', 'telegramApiHash', 'googleDriveAuth', 'googleDriveUploadFolderIds'] self.optVars: typing.List[str] = ['ariaGlobalOpts', 'authorizedChats', 'dlRootDir', 'logLevel', - 'megaAuth', 'statusUpdateInterval', 'trackersListUrl'] + 'megaAuth', 'statusUpdateInterval', 'trackersListUrl', 'ytdlFormat'] self.optVals: typing.List[typing.Union[str, typing.Dict]] = \ [{'allow-overwrite': 'true', 'bt-max-peers': '0', 'follow-torrent': 'mem', 'max-connection-per-server': '8', 'max-overall-upload-limit': '1K', 'min-split-size': '10M', 'seed-time': '0.01', 'split': '10'}, - {}, 'dl', 'INFO', {}, '5', 'https://trackerslist.com/all_aria2.txt'] + {}, 'dl', 'INFO', {}, '5', 'https://trackerslist.com/all_aria2.txt', 'best/bestvideo+bestaudio'] self.emptyVals: typing.List[typing.Union[str, typing.Dict]] = ['', ' ', {}] self.isFixConfigJson: bool = False self.configVarsLoad() @@ -389,6 +387,14 @@ def fileHash(filePath: str) -> str: fileChunk = fileStream.read(blockSize) return hashSum.hexdigest() + @staticmethod + def folderSize(folderPath: str) -> int: + size: int = 0 + for path, dirs, files in os.walk(folderPath): + for file in files: + size += os.path.getsize(os.path.join(path, file)) + return size + def progressBar(self, progress: float) -> str: progressRounded = round(progress) numFull = progressRounded // 8 @@ -1034,7 +1040,7 @@ def stageFive(self, update: telegram.Update, _: telegram.ext.CallbackContext) -> query.edit_message_text(text='addMirror Cancelled !') return telegram.ext.ConversationHandler.END - # TODO: reduce this function code if possible + # TODO: reduce this method code if possible def getMirrorInfoStr(self): mirrorInfoStr = f'[uid | {self.mirrorInfo.uid}]\n' if self.mirrorInfo.isAriaDownload: @@ -1075,7 +1081,7 @@ def initHelper(self) -> None: def addMirror(self, mirrorInfo: 'MirrorInfo') -> None: self.logger.debug(vars(mirrorInfo)) self.mirrorInfos[mirrorInfo.uid] = mirrorInfo - self.mirrorInfos[mirrorInfo.uid].timeStart = int(time.time()) + self.mirrorInfos[mirrorInfo.uid].timeStart = time.time() self.botHelper.mirrorListenerHelper.updateStatus(mirrorInfo.uid, MirrorStatus.addMirror) self.botHelper.threadingHelper.initThread(target=self.botHelper.statusHelper.addStatus, name=f'{mirrorInfo.uid}-addStatus', chatId=mirrorInfo.chatId, msgId=mirrorInfo.msgId) @@ -1231,7 +1237,7 @@ def updateProgress(self, uid: str) -> None: {MirrorInfo.updatableVars[0]: dlObj.total_length, MirrorInfo.updatableVars[1]: dlObj.completed_length, MirrorInfo.updatableVars[2]: dlObj.download_speed, - MirrorInfo.updatableVars[3]: int(time.time())} + MirrorInfo.updatableVars[3]: time.time()} if dlObj.is_torrent: currVars[MirrorInfo.updatableVars[4]] = True currVars[MirrorInfo.updatableVars[5]] = dlObj.num_seeders @@ -1294,16 +1300,16 @@ def addDownload(self, mirrorInfo: 'MirrorInfo') -> None: isFolder = True if mirrorInfo.isGoogleDriveUpload and not (mirrorInfo.isCompress or mirrorInfo.isDecompress): if isFolder: - folderId = self.cloneFolder(sourceFolderId=sourceId, parentFolderId=mirrorInfo.googleDriveUploadFolderId) + folderId = self.cloneFolder(sourceFolderId=sourceId, parentFolderId=mirrorInfo.googleDriveUploadFolderId, uid=mirrorInfo.uid) self.botHelper.mirrorHelper.mirrorInfos[mirrorInfo.uid].uploadUrl = self.baseFolderDownloadUrl.format(folderId) else: - fileId = self.cloneFile(sourceFileId=sourceId, parentFolderId=mirrorInfo.googleDriveUploadFolderId) + fileId = self.cloneFile(sourceFileId=sourceId, parentFolderId=mirrorInfo.googleDriveUploadFolderId, uid=mirrorInfo.uid) self.botHelper.mirrorHelper.mirrorInfos[mirrorInfo.uid].uploadUrl = self.baseFileDownloadUrl.format(fileId) else: if isFolder: - self.downloadFolder(sourceFolderId=sourceId, dlPath=mirrorInfo.path) + self.downloadFolder(sourceFolderId=sourceId, dlPath=mirrorInfo.path, uid=mirrorInfo.uid) else: - self.downloadFile(sourceFileId=sourceId, dlPath=mirrorInfo.path) + self.downloadFile(sourceFileId=sourceId, dlPath=mirrorInfo.path, uid=mirrorInfo.uid) self.botHelper.mirrorListenerHelper.updateStatus(mirrorInfo.uid, MirrorStatus.downloadComplete) def cancelDownload(self, uid: str) -> None: @@ -1311,12 +1317,14 @@ def cancelDownload(self, uid: str) -> None: def addUpload(self, mirrorInfo: 'MirrorInfo') -> None: if not (mirrorInfo.isGoogleDriveDownload and not (mirrorInfo.isCompress or mirrorInfo.isDecompress)): + currVars = {MirrorInfo.updatableVars[0]: self.botHelper.getHelper.folderSize(mirrorInfo.path)} + self.botHelper.mirrorHelper.mirrorInfos[mirrorInfo.uid].updateVars(currVars) uploadPath = os.path.join(mirrorInfo.path, os.listdir(mirrorInfo.path)[0]) if os.path.isdir(uploadPath): - folderId = self.uploadFolder(folderPath=uploadPath, parentFolderId=mirrorInfo.googleDriveUploadFolderId) + folderId = self.uploadFolder(folderPath=uploadPath, parentFolderId=mirrorInfo.googleDriveUploadFolderId, uid=mirrorInfo.uid) self.botHelper.mirrorHelper.mirrorInfos[mirrorInfo.uid].uploadUrl = self.baseFolderDownloadUrl.format(folderId) if os.path.isfile(uploadPath): - fileId = self.uploadFile(filePath=uploadPath, parentFolderId=mirrorInfo.googleDriveUploadFolderId) + fileId = self.uploadFile(filePath=uploadPath, parentFolderId=mirrorInfo.googleDriveUploadFolderId, uid=mirrorInfo.uid) self.botHelper.mirrorHelper.mirrorInfos[mirrorInfo.uid].uploadUrl = self.baseFileDownloadUrl.format(fileId) else: time.sleep(self.botHelper.statusHelper.statusUpdateInterval) @@ -1339,12 +1347,14 @@ def authorizeApi(self) -> None: else: self.logger.info('Google Drive API User Token Needs to Refreshed Manually ! Exiting...') exit(1) + else: + self.buildService() def buildService(self) -> None: self.service = googleapiclient.discovery.build(serviceName='drive', version='v3', credentials=self.oauthCreds, cache_discovery=False) - def uploadFile(self, filePath: str, parentFolderId: str) -> str: + def uploadFile(self, filePath: str, parentFolderId: str, uid: str) -> str: upStatus: googleapiclient.http.MediaUploadProgress fileName, fileMimeType, fileMetadata, mediaBody = self.getUpData(filePath, isResumable=True) fileMetadata['parents'] = [parentFolderId] @@ -1352,9 +1362,11 @@ def uploadFile(self, filePath: str, parentFolderId: str) -> str: upResponse = None while not upResponse: upStatus, upResponse = fileOp.next_chunk() + sizeUpdate = (self.chunkSize if not upResponse else (os.path.getsize(filePath) % self.chunkSize)) + self.updateProgress(sizeUpdate, uid) return upResponse['id'] - def uploadFolder(self, folderPath: str, parentFolderId: str) -> str: + def uploadFolder(self, folderPath: str, parentFolderId: str, uid: str) -> str: folderName = folderPath.split('/')[-1] folderId = self.createFolder(folderName, parentFolderId) folderContents = os.listdir(folderPath) @@ -1362,28 +1374,30 @@ def uploadFolder(self, folderPath: str, parentFolderId: str) -> str: for contentName in folderContents: contentPath = os.path.join(folderPath, contentName) if os.path.isdir(contentPath): - self.uploadFolder(contentPath, folderId) + self.uploadFolder(folderPath=contentPath, parentFolderId=folderId, uid=uid) if os.path.isfile(contentPath): - self.uploadFile(contentPath, folderId) + self.uploadFile(filePath=contentPath, parentFolderId=folderId, uid=uid) return folderId - def cloneFile(self, sourceFileId: str, parentFolderId: str) -> str: + def cloneFile(self, sourceFileId: str, parentFolderId: str, uid: str) -> str: fileMetadata = {'parents': [parentFolderId]} - return self.service.files().copy(supportsAllDrives=True, fileId=sourceFileId, body=fileMetadata).execute()['id'] + fileOp = self.service.files().copy(supportsAllDrives=True, fileId=sourceFileId, body=fileMetadata).execute() + self.updateProgress(self.getSizeById(sourceFileId), uid) + return fileOp['id'] - def cloneFolder(self, sourceFolderId: str, parentFolderId: str) -> str: + def cloneFolder(self, sourceFolderId: str, parentFolderId: str, uid: str) -> str: sourceFolderName = self.getMetadataById(sourceFolderId, 'name') folderId = self.createFolder(sourceFolderName, parentFolderId) folderContents = self.getFolderContentsById(sourceFolderId) if len(folderContents) != 0: for content in folderContents: if content.get('mimeType') == self.googleDriveFolderMimeType: - self.cloneFolder(content.get('id'), folderId) + self.cloneFolder(sourceFolderId=content.get('id'), parentFolderId=folderId, uid=uid) else: - self.cloneFile(content.get('id'), folderId) + self.cloneFile(sourceFileId=content.get('id'), parentFolderId=folderId, uid=uid) return folderId - def downloadFile(self, sourceFileId: str, dlPath: str) -> None: + def downloadFile(self, sourceFileId: str, dlPath: str, uid: str) -> None: fileName = self.getMetadataById(sourceFileId, 'name') filePath = os.path.join(dlPath, fileName) downStatus: googleapiclient.http.MediaDownloadProgress @@ -1393,9 +1407,11 @@ def downloadFile(self, sourceFileId: str, dlPath: str) -> None: downResponse = None while not downResponse: downStatus, downResponse = fileOp.next_chunk() + sizeUpdate = (self.chunkSize if not downResponse else (self.getSizeById(sourceFileId) % self.chunkSize)) + self.updateProgress(sizeUpdate, uid) return - def downloadFolder(self, sourceFolderId: str, dlPath: str) -> None: + def downloadFolder(self, sourceFolderId: str, dlPath: str, uid: str) -> None: folderName = self.getMetadataById(sourceFolderId, 'name') folderPath = os.path.join(dlPath, folderName) os.mkdir(folderPath) @@ -1403,9 +1419,9 @@ def downloadFolder(self, sourceFolderId: str, dlPath: str) -> None: if len(folderContents) != 0: for content in folderContents: if content.get('mimeType') == self.googleDriveFolderMimeType: - self.downloadFolder(content.get('id'), folderPath) + self.downloadFolder(sourceFolderId=content.get('id'), dlPath=folderPath, uid=uid) else: - self.downloadFile(content.get('id'), folderPath) + self.downloadFile(sourceFileId=content.get('id'), dlPath=folderPath, uid=uid) return def createFolder(self, folderName: str, parentFolderId: str) -> str: @@ -1475,6 +1491,18 @@ def patchFile(self, filePath: str, fileId: str) -> str: fileOp = self.service.files().update(fileId=fileId, body=fileMetadata, media_body=mediaBody).execute() return f"Patched: [{fileOp['id']}] [{fileName}] [{os.path.getsize(fileName)} bytes]" + def updateProgress(self, sizeUpdate: int, uid: str): + sizeLast = self.botHelper.mirrorHelper.mirrorInfos[uid].sizeCurrent + timeLast = self.botHelper.mirrorHelper.mirrorInfos[uid].timeCurrent + speedLast = self.botHelper.mirrorHelper.mirrorInfos[uid].speedCurrent + sizeCurrent = sizeLast + sizeUpdate + timeCurrent = time.time() + timeDiff = timeCurrent - timeLast + speedCurrent = (int(sizeUpdate / timeDiff) if timeDiff else speedLast) + self.botHelper.mirrorHelper.mirrorInfos[uid].updateVars({MirrorInfo.updatableVars[1]: sizeCurrent, + MirrorInfo.updatableVars[2]: speedCurrent, + MirrorInfo.updatableVars[3]: timeCurrent}) + class MegaHelper(BaseHelper): def __init__(self, botHelper: BotHelper): @@ -1570,6 +1598,8 @@ def cancelDownload(self, uid: str) -> None: raise NotImplementedError def addUpload(self, mirrorInfo: 'MirrorInfo') -> None: + currVars = {MirrorInfo.updatableVars[0]: self.botHelper.getHelper.folderSize(mirrorInfo.path)} + self.botHelper.mirrorHelper.mirrorInfos[mirrorInfo.uid].updateVars(currVars) uploadPath = os.path.join(mirrorInfo.path, os.listdir(mirrorInfo.path)[0]) upResponse: bool = True if os.path.isfile(uploadPath): @@ -1627,10 +1657,9 @@ def initHelper(self) -> None: super().initHelper() def addDownload(self, mirrorInfo: 'MirrorInfo') -> None: - ytdlOpts: dict = {'format': 'best/bestvideo+bestaudio', 'logger': self.logger, + ytdlOpts: dict = {'quiet': True, 'format': mirrorInfo.ytdlFormat, 'progress_hooks': [self.progressHook], 'outtmpl': f'{mirrorInfo.path}/%(title)s-%(id)s.f%(format_id)s.%(ext)s'} self.downloadVideo(mirrorInfo.downloadUrl, ytdlOpts) - self.botHelper.mirrorListenerHelper.updateStatus(mirrorInfo.uid, MirrorStatus.downloadComplete) def cancelDownload(self, uid: str) -> None: raise NotImplementedError @@ -1640,6 +1669,18 @@ def downloadVideo(videoUrl: str, ytdlOpts: dict) -> None: with youtube_dl.YoutubeDL(ytdlOpts) as ytdl: ytdl.download([videoUrl]) + def progressHook(self, progressUpdate: dict): + uid = progressUpdate['filename'].replace(self.botHelper.envVars['dlRootDirPath'], '').split('/')[1] + if progressUpdate['status'] == 'downloading': + currVars: typing.Dict[str, typing.Union[int, float, str]] = \ + {MirrorInfo.updatableVars[0]: int((sizeTotal if (sizeTotal := progressUpdate['total_bytes']) else 0)), + MirrorInfo.updatableVars[1]: int((sizeCurrent if (sizeCurrent := progressUpdate['downloaded_bytes']) else 0)), + MirrorInfo.updatableVars[2]: int((speedCurrent if (speedCurrent := progressUpdate['speed']) else 0)), + MirrorInfo.updatableVars[3]: time.time()} + self.botHelper.mirrorHelper.mirrorInfos[uid].updateVars(currVars) + if progressUpdate['status'] == 'finished': + self.botHelper.mirrorListenerHelper.updateStatus(uid, MirrorStatus.downloadComplete) + class CompressionHelper(BaseHelper): def __init__(self, botHelper: BotHelper): @@ -1734,8 +1775,8 @@ def getStatusMsgTxt(self) -> str: for uid in self.botHelper.mirrorHelper.mirrorInfos.keys(): mirrorInfo: MirrorInfo = self.botHelper.mirrorHelper.mirrorInfos[uid] statusMsgTxt += f'{mirrorInfo.uid} | {mirrorInfo.status}\n' - if mirrorInfo.status == MirrorStatus.downloadProgress: - if mirrorInfo.isAriaDownload: + if mirrorInfo.status in [MirrorStatus.downloadProgress, MirrorStatus.uploadProgress]: + if mirrorInfo.status == MirrorStatus.downloadProgress and mirrorInfo.isAriaDownload: self.botHelper.ariaHelper.updateProgress(mirrorInfo.uid) statusMsgTxt += f'S: {self.botHelper.getHelper.readableSize(mirrorInfo.sizeCurrent)} | ' \ f'{self.botHelper.getHelper.readableSize(mirrorInfo.sizeTotal)} | ' \ @@ -1743,8 +1784,8 @@ def getStatusMsgTxt(self) -> str: f'P: {self.botHelper.getHelper.progressBar(mirrorInfo.progressPercent)} | ' \ f'{mirrorInfo.progressPercent}% | ' \ f'{self.botHelper.getHelper.readableSize(mirrorInfo.speedCurrent)}/s\n' \ - f'T: {self.botHelper.getHelper.readableTime(mirrorInfo.timeCurrent - mirrorInfo.timeStart)} | ' \ - f'{self.botHelper.getHelper.readableTime(mirrorInfo.timeEnd - mirrorInfo.timeCurrent)}\n' + f'T: {self.botHelper.getHelper.readableTime(int(mirrorInfo.timeCurrent - mirrorInfo.timeStart))} | ' \ + f'{self.botHelper.getHelper.readableTime(int(mirrorInfo.timeEnd - mirrorInfo.timeCurrent))}\n' statusMsgTxt += (f'nS: {mirrorInfo.numSeeders} nL: {mirrorInfo.numLeechers}\n' if mirrorInfo.isTorrent else '') return statusMsgTxt @@ -2016,8 +2057,7 @@ def onUploadError(self, mirrorInfo: 'MirrorInfo') -> None: self.checkUploadQueue() def resetMirrorProgress(self, uid: str) -> None: - self.botHelper.mirrorHelper.mirrorInfos[uid].timeEnd = 0 - self.botHelper.mirrorHelper.mirrorInfos[uid].progressPercent = 0 + self.botHelper.mirrorHelper.mirrorInfos[uid].resetVars() class MegaApiWrapper: @@ -2127,7 +2167,7 @@ def onTransferUpdate(self, api: mega.MegaApi, transfer: mega.MegaTransfer): {MirrorInfo.updatableVars[0]: transfer.getTotalBytes(), MirrorInfo.updatableVars[1]: transfer.getTransferredBytes(), MirrorInfo.updatableVars[2]: transfer.getSpeed(), - MirrorInfo.updatableVars[3]: int(time.time())} + MirrorInfo.updatableVars[3]: time.time()} self.megaHelper.botHelper.mirrorHelper.mirrorInfos[uid].updateVars(currVars) self.logger.debug(f'Transfer Update ({transfer} {transfer.getFileName()}); ' f'Progress: {transfer.getTransferredBytes() / 1024} KB of {transfer.getTotalBytes() / 1024} KB, ' @@ -2157,14 +2197,15 @@ def __init__(self, msg: telegram.Message, botHelper: BotHelper): self.uid: str = botHelper.getHelper.randomString(8) self.path: str = os.path.join(botHelper.envVars['dlRootDirPath'], self.uid) self.status: str = '' - self.downloadUrl: str = '' self.tag: str = '' - self.timeStart: int = 0 - self.timeCurrent: int = 0 - self.timeEnd: int = 0 - self.speedCurrent: int = 0 + self.downloadUrl: str = '' + self.ytdlFormat: str = botHelper.configHelper.configVars[botHelper.configHelper.optVars[7]] self.sizeTotal: int = 0 self.sizeCurrent: int = 0 + self.timeStart: float = 0.0 + self.timeCurrent: float = 0.0 + self.timeEnd: float = 0.0 + self.speedCurrent: int = 0 self.progressPercent: float = 0.0 self.isTorrent: bool = False self.numSeeders: int = 0 @@ -2182,6 +2223,12 @@ def __init__(self, msg: telegram.Message, botHelper: BotHelper): self.isCompress: bool = False self.isDecompress: bool = False + def resetVars(self): + self.sizeTotal, self.sizeCurrent = 0, 0 + self.timeEnd, self.timeCurrent = 0.0, 0.0 + self.speedCurrent = 0 + self.progressPercent = 0.0 + def updateVars(self, currVars: typing.Dict[str, typing.Union[int, float, str]]) -> None: currVarsKeys = list(currVars.keys()) if self.updatableVars[0] in currVarsKeys: @@ -2193,7 +2240,7 @@ def updateVars(self, currVars: typing.Dict[str, typing.Union[int, float, str]]) if self.sizeTotal != 0: self.progressPercent = round(((self.sizeCurrent / self.sizeTotal) * 100), ndigits=2) if self.speedCurrent != 0: - self.timeEnd = self.timeCurrent + int((self.sizeTotal - self.sizeCurrent) / self.speedCurrent) + self.timeEnd = self.timeCurrent + ((self.sizeTotal - self.sizeCurrent) / self.speedCurrent) if self.updatableVars[4] in currVarsKeys: self.isTorrent = True self.numSeeders = currVars[self.updatableVars[5]]