分模块
This commit is contained in:
limil
parent
b84358f228
commit
2aa4d0870b
352
PikPakFileSystem.py
Normal file
352
PikPakFileSystem.py
Normal file
@ -0,0 +1,352 @@
|
||||
import httpx
|
||||
from pikpakapi import PikPakApi, DownloadStatus
|
||||
from typing import Dict
|
||||
from datetime import datetime
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
class NodeBase:
|
||||
def __init__(self, id : str, name : str, fatherId : str):
|
||||
self.id = id
|
||||
self.name = name
|
||||
self._father_id = fatherId
|
||||
self.lastUpdate : datetime = None
|
||||
|
||||
class DirNode(NodeBase):
|
||||
def __init__(self, id : str, name : str, fatherId : str):
|
||||
super().__init__(id, name, fatherId)
|
||||
self.children_id : list[str] = []
|
||||
|
||||
class FileNode(NodeBase):
|
||||
def __init__(self, id : str, name : str, fatherId : str):
|
||||
super().__init__(id, name, fatherId)
|
||||
self.url : str = None
|
||||
|
||||
class PikPakFileSystem:
|
||||
#region 内部接口
|
||||
def __init__(self, auth_cache_path : str = None, proxy_address : str = None, root_id : str = None):
|
||||
# 初始化虚拟文件节点
|
||||
self._nodes : Dict[str, NodeBase] = {}
|
||||
self._root : DirNode = DirNode(root_id, "", None)
|
||||
self._cwd : DirNode = self._root
|
||||
|
||||
# 初始化鉴权和代理信息
|
||||
self._auth_cache_path : str = auth_cache_path
|
||||
self.proxy_address : str = proxy_address
|
||||
self._pikpak_client : PikPakApi = None
|
||||
self._try_login_from_cache()
|
||||
|
||||
|
||||
#region 鉴权信息相关
|
||||
class PikPakToken:
|
||||
def __init__(self, username : str, password : str, access_token : str, refresh_token : str, user_id : str):
|
||||
self.username : str = username
|
||||
self.password : str = password
|
||||
self.access_token : str = access_token
|
||||
self.refresh_token : str = refresh_token
|
||||
self.user_id : str = user_id
|
||||
|
||||
def to_json(self):
|
||||
return json.dumps(self.__dict__)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_str):
|
||||
data = json.loads(json_str)
|
||||
return cls(**data)
|
||||
|
||||
def _init_client_by_token(self, token : PikPakToken) -> None:
|
||||
self._init_client_by_username_and_password(token.username, token.password)
|
||||
self._pikpak_client.access_token = token.access_token
|
||||
self._pikpak_client.refresh_token = token.refresh_token
|
||||
self._pikpak_client.user_id = token.user_id
|
||||
self._pikpak_client.encode_token()
|
||||
|
||||
def _init_client_by_username_and_password(self, username : str, password : str) -> None:
|
||||
httpx_client_args : Dict[str, Any] = None
|
||||
if self.proxy_address != None:
|
||||
httpx_client_args = {
|
||||
"proxy": self.proxy_address,
|
||||
"transport": httpx.AsyncHTTPTransport()
|
||||
}
|
||||
|
||||
self._pikpak_client = PikPakApi(
|
||||
username = username,
|
||||
password = password,
|
||||
httpx_client_args=httpx_client_args)
|
||||
|
||||
def _try_login_from_cache(self) -> None:
|
||||
if self._auth_cache_path is None:
|
||||
return
|
||||
if not os.path.exists(self._auth_cache_path):
|
||||
return
|
||||
try:
|
||||
with open(self._auth_cache_path, 'r', encoding='utf-8') as file:
|
||||
content : str = file.read()
|
||||
token : PikPakFileSystem.PikPakToken = PikPakFileSystem.PikPakToken.from_json(content)
|
||||
self._init_client_by_token(token)
|
||||
logging.info("successfully load login info from cache")
|
||||
except Exception as e:
|
||||
logging.error(f"failed to load login info from cache, exception occurred: {e}")
|
||||
|
||||
def _dump_login_info(self) -> None:
|
||||
if self._auth_cache_path is None:
|
||||
return
|
||||
with open(self._auth_cache_path, 'w', encoding='utf-8') as file:
|
||||
token : PikPakFileSystem.PikPakToken = PikPakFileSystem.PikPakToken(self._pikpak_client.username, self._pikpak_client.password, self._pikpak_client.access_token, self._pikpak_client.refresh_token, self._pikpak_client.user_id)
|
||||
file.write(token.to_json())
|
||||
logging.info("successfully dump login info to cache")
|
||||
|
||||
#endregion
|
||||
|
||||
#region 文件系统相关
|
||||
class PathWalker():
|
||||
def __init__(self, path : str, sep : str = "/"):
|
||||
self._path_spots : list[str] = []
|
||||
if not path.startswith(sep):
|
||||
self._path_spots.append(".")
|
||||
path_spots : list[str] = path.split(sep)
|
||||
self._path_spots.extend(path_spots)
|
||||
|
||||
def IsAbsolute(self) -> bool:
|
||||
return len(self._path_spots) == 0 or self._path_spots[0] != "."
|
||||
|
||||
def AppendSpot(self, spot) -> None:
|
||||
self._path_spots.append(spot)
|
||||
|
||||
def Walk(self) -> list[str]:
|
||||
return self._path_spots
|
||||
|
||||
async def _get_node_by_id(self, id : str) -> NodeBase:
|
||||
if id == self._root.id:
|
||||
return self._root
|
||||
if id not in self._nodes:
|
||||
return None
|
||||
return self._nodes[id]
|
||||
|
||||
async def _get_father_node(self, node : NodeBase) -> NodeBase:
|
||||
if node is self._root:
|
||||
return self._root
|
||||
return await self._get_node_by_id(node._father_id)
|
||||
|
||||
async def _add_node(self, node : NodeBase) -> None:
|
||||
self._nodes[node.id] = node
|
||||
father = await self._get_father_node(node)
|
||||
if father is not None and isinstance(father, DirNode):
|
||||
father.children_id.append(node.id)
|
||||
|
||||
async def _remove_node(self, node : NodeBase) -> None:
|
||||
father = await self._get_father_node(node)
|
||||
if father is not None and isinstance(father, DirNode):
|
||||
father.children_id.remove(node.id)
|
||||
self._nodes.pop(node.id)
|
||||
|
||||
async def _find_child_in_dir_by_name(self, dir : DirNode, name : str) -> NodeBase:
|
||||
if dir is self._root and name == "":
|
||||
return self._root
|
||||
for child_id in dir.children_id:
|
||||
node = await self._get_node_by_id(child_id)
|
||||
if node.name == name:
|
||||
return node
|
||||
return None
|
||||
|
||||
async def _refresh(self, node : NodeBase):
|
||||
if isinstance(node, DirNode):
|
||||
if node.lastUpdate != None:
|
||||
return
|
||||
next_page_token : str = None
|
||||
children_info : list[Dict[str, Any]] = []
|
||||
while True:
|
||||
dir_info : Dict[str, Any] = await self._pikpak_client.file_list(parent_id = node.id, next_page_token=next_page_token)
|
||||
next_page_token = dir_info["next_page_token"]
|
||||
children_info.extend(dir_info["files"])
|
||||
if next_page_token is None or next_page_token == "":
|
||||
break
|
||||
|
||||
node.children_id.clear()
|
||||
for child_info in children_info:
|
||||
id : str = child_info["id"]
|
||||
name : str = child_info["name"]
|
||||
|
||||
child : NodeBase = await self._get_node_by_id(id)
|
||||
if child is None:
|
||||
if child_info["kind"].endswith("folder"):
|
||||
child = DirNode(id, name, node.id)
|
||||
else:
|
||||
child = FileNode(id, name, node.id)
|
||||
child.name = name
|
||||
child._father_id = node.id
|
||||
await self._add_node(child)
|
||||
elif isinstance(node, FileNode):
|
||||
result = await self._pikpak_client.get_download_url(node.id)
|
||||
node.url = result["web_content_link"]
|
||||
|
||||
node.lastUpdate = datetime.now()
|
||||
|
||||
async def _path_to_node(self, path : str) -> NodeBase:
|
||||
father, son_name = await self._path_to_father_node_and_son_name(path)
|
||||
if son_name == "":
|
||||
return father
|
||||
if isinstance(father, DirNode):
|
||||
return await self._find_child_in_dir_by_name(father, son_name)
|
||||
return None
|
||||
|
||||
async def _path_to_father_node_and_son_name(self, path : str) -> tuple[NodeBase, str]:
|
||||
path_walker : PikPakFileSystem.PathWalker = PikPakFileSystem.PathWalker(path)
|
||||
father : NodeBase = None
|
||||
son_name : str = None
|
||||
current : NodeBase = self._root if path_walker.IsAbsolute() else self._cwd
|
||||
|
||||
for spot in path_walker.Walk():
|
||||
if current is None:
|
||||
father = None
|
||||
break
|
||||
if spot == "..":
|
||||
current = await self._get_father_node(current)
|
||||
continue
|
||||
father = current
|
||||
if not isinstance(current, DirNode):
|
||||
current = None
|
||||
continue
|
||||
await self._refresh(current)
|
||||
if spot == ".":
|
||||
continue
|
||||
sonName = spot
|
||||
current = await self._find_child_in_dir_by_name(current, spot)
|
||||
|
||||
if current != None:
|
||||
father = await self._get_father_node(current)
|
||||
sonName = current.name
|
||||
|
||||
return father, sonName
|
||||
|
||||
async def _node_to_path(self, node : NodeBase) -> str:
|
||||
if node is self._root:
|
||||
return "/"
|
||||
spots : list[str] = []
|
||||
current = node
|
||||
while current is not self._root:
|
||||
spots.append(current.name)
|
||||
current = await self._get_father_node(current)
|
||||
spots.append("")
|
||||
return "/".join(reversed(spots))
|
||||
|
||||
async def _is_ancestors_of(self, node_a : NodeBase, node_b : NodeBase) -> bool:
|
||||
if node_b is node_a:
|
||||
return False
|
||||
if node_a is self._root:
|
||||
return True
|
||||
while node_b._father_id != self._root.id:
|
||||
node_b = await self._get_father_node(node_b)
|
||||
if node_b is node_a:
|
||||
return True
|
||||
return False
|
||||
#endregion
|
||||
|
||||
#endregion
|
||||
|
||||
#region 对外接口
|
||||
async def Login(self, username : str = None, password : str = None) -> None:
|
||||
if self._pikpak_client != None and username is None and password is None:
|
||||
username = self._pikpak_client.username
|
||||
password = self._pikpak_client.password
|
||||
|
||||
if username == None and password == None:
|
||||
raise Exception("Username and password are required")
|
||||
|
||||
self._init_client_by_username_and_password(username, password)
|
||||
await self._pikpak_client.login()
|
||||
self._dump_login_info()
|
||||
|
||||
async def IsDir(self, path : str) -> bool:
|
||||
node = await self._path_to_node(path)
|
||||
return isinstance(node, DirNode)
|
||||
|
||||
async def SplitPath(self, path : str) -> tuple[str, str]:
|
||||
father, son_name = await self._path_to_father_node_and_son_name(path)
|
||||
return await self._node_to_path(father), son_name
|
||||
|
||||
async def GetChildrenNames(self, path : str, ignore_files : bool) -> list[str]:
|
||||
node = await self._path_to_node(path)
|
||||
if not isinstance(node, DirNode):
|
||||
return []
|
||||
await self._refresh(node)
|
||||
children_names : list[str] = []
|
||||
for child_id in node.children_id:
|
||||
child = await self._get_node_by_id(child_id)
|
||||
if ignore_files and isinstance(child, FileNode):
|
||||
continue
|
||||
children_names.append(child.name)
|
||||
return children_names
|
||||
|
||||
async def Delete(self, paths : list[str]) -> None:
|
||||
nodes = [await self._path_to_node(path) for path in paths]
|
||||
for node in nodes:
|
||||
if await self._is_ancestors_of(node, self._cwd):
|
||||
raise Exception("Cannot delete ancestors")
|
||||
await self._pikpak_client.delete_to_trash([node.id for node in nodes])
|
||||
for node in nodes:
|
||||
await self._remove_node(node)
|
||||
|
||||
async def MakeDir(self, path : str) -> None:
|
||||
father_path, son_name = await self._path_to_father_node_and_son_name(path)
|
||||
father = await self._path_to_node(father_path)
|
||||
result = await self._pikpak_client.create_folder(son_name, father.id)
|
||||
id = result["file"]["id"]
|
||||
name = result["file"]["name"]
|
||||
son = DirNode(id, name, father.id)
|
||||
await self._add_node(son)
|
||||
|
||||
async def SetCwd(self, path : str) -> None:
|
||||
node = await self._path_to_node(path)
|
||||
if not isinstance(node, DirNode):
|
||||
raise Exception("Not a directory")
|
||||
self._cwd = node
|
||||
|
||||
async def GetCwd(self) -> str:
|
||||
return await self._node_to_path(self._cwd)
|
||||
|
||||
async def GetNodeIdByPath(self, path : str) -> str:
|
||||
node = await self._path_to_node(path)
|
||||
if node is None:
|
||||
return None
|
||||
return node.id
|
||||
|
||||
async def IfExists(self, path : str) -> bool:
|
||||
return await self._path_to_node(path) is not None
|
||||
|
||||
async def ToFullPath(self, path : str) -> str:
|
||||
node = await self._path_to_node(path)
|
||||
return await self._node_to_path(node)
|
||||
|
||||
async def RemoteDownload(self, torrent : str, remote_base_path : str) -> tuple[str, str]:
|
||||
node = await self._path_to_node(remote_base_path)
|
||||
info = await self._pikpak_client.offline_download(torrent, node.id)
|
||||
return info["task"]["file_id"], info["task"]["id"]
|
||||
|
||||
async def QueryTaskStatus(self, task_id : str, node_id : str) -> DownloadStatus:
|
||||
return await self._pikpak_client.get_task_status(task_id, node_id)
|
||||
|
||||
async def UpdateDirectory(self, path : str, son_id : str) -> str:
|
||||
await self._path_to_node(path)
|
||||
son_info = await self._pikpak_client.offline_file_info(son_id)
|
||||
kind = son_info["kind"]
|
||||
parent_id = son_info["parent_id"]
|
||||
name = son_info["name"]
|
||||
son : NodeBase = None
|
||||
if kind.endswith("folder"):
|
||||
son = DirNode(son_id, name, parent_id)
|
||||
else:
|
||||
son = FileNode(son_id, name, parent_id)
|
||||
await self._add_node(son)
|
||||
return await self._node_to_path(son)
|
||||
|
||||
async def JoinPath(self, father : str, son : str) -> str:
|
||||
father_node = await self._path_to_node(father)
|
||||
son_node = await self._find_child_in_dir_by_name(father_node, son)
|
||||
if son_node is None:
|
||||
raise Exception("Son not found")
|
||||
return await self._node_to_path(son_node)
|
||||
|
||||
#endregion
|
||||
508
PikPakFs.py
508
PikPakFs.py
@ -1,508 +0,0 @@
|
||||
import httpx
|
||||
from pikpakapi import PikPakApi, DownloadStatus
|
||||
from typing import Dict
|
||||
from datetime import datetime
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from enum import Enum
|
||||
import asyncio
|
||||
import shortuuid
|
||||
from utils import PathWalker
|
||||
from typing import Callable, Awaitable
|
||||
import random
|
||||
|
||||
class TaskStatus(Enum):
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
DONE = "done"
|
||||
ERROR = "error"
|
||||
PAUSED = "paused"
|
||||
|
||||
class PikPakTaskStatus(Enum):
|
||||
PENDING = "pending"
|
||||
REMOTE_DOWNLOADING = "remote"
|
||||
LOCAL_DOWNLOADING = "local"
|
||||
DONE = "done"
|
||||
|
||||
class FileDownloadTaskStatus(Enum):
|
||||
PENDING = "pending"
|
||||
DOWNLOADING = "downloading"
|
||||
DONE = "done"
|
||||
|
||||
class UnRecoverableError(Exception):
|
||||
def __init__(self, message):
|
||||
super().__init__(message)
|
||||
|
||||
class TaskBase:
|
||||
def __init__(self, id : str, tag : str = "", maxConcurrentNumber = -1):
|
||||
self.id : str = shortuuid.uuid() if id is None else id
|
||||
self.tag : str = tag
|
||||
self.maxConcurrentNumber : int = maxConcurrentNumber
|
||||
self.name : str = ""
|
||||
|
||||
self._status : TaskStatus = TaskStatus.PENDING
|
||||
|
||||
self.worker : asyncio.Task = None
|
||||
self.handler : Callable[..., Awaitable] = None
|
||||
|
||||
class PikPakTask(TaskBase):
|
||||
TAG = "PikPakTask"
|
||||
MAX_CONCURRENT_NUMBER = 5
|
||||
|
||||
def __init__(self, torrent : str, toDirId : str, nodeId : str = None, status : PikPakTaskStatus = PikPakTaskStatus.PENDING, id : str = None):
|
||||
super().__init__(id, PikPakTask.TAG, PikPakTask.MAX_CONCURRENT_NUMBER)
|
||||
self.status : PikPakTaskStatus = status
|
||||
self.toDirId : str = toDirId
|
||||
self.nodeId : str = nodeId
|
||||
self.torrent : str = torrent # todo: 将torrent的附加参数去掉再加入
|
||||
self.remoteTaskId : str = None
|
||||
|
||||
self.progress : str = ""
|
||||
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, PikPakTask):
|
||||
return self is other or self.nodeId == other.nodeId
|
||||
return False
|
||||
|
||||
class FileDownloadTask(TaskBase):
|
||||
TAG = "FileDownloadTask"
|
||||
MAX_CONCURRENT_NUMBER = 5
|
||||
|
||||
def __init__(self, nodeId : str, PikPakTaskId : str, relativePath : str, status : FileDownloadTaskStatus = FileDownloadTaskStatus.PENDING, id : str = None):
|
||||
super().__init__(id, FileDownloadTask.TAG, FileDownloadTask.MAX_CONCURRENT_NUMBER)
|
||||
self.status : FileDownloadTaskStatus = status
|
||||
self.PikPakTaskId : str = PikPakTaskId
|
||||
self.nodeId : str = nodeId
|
||||
self.relativePath : str = relativePath
|
||||
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, FileDownloadTask):
|
||||
return self is other or (self.nodeId == other.nodeId and self.PikPakTaskId == other.PikPakTaskId)
|
||||
return False
|
||||
|
||||
async def TaskWorker(task : TaskBase):
|
||||
try:
|
||||
if task._status != TaskStatus.PENDING:
|
||||
return
|
||||
task._status = TaskStatus.RUNNING
|
||||
await task.handler(task)
|
||||
task._status = TaskStatus.DONE
|
||||
except asyncio.CancelledError:
|
||||
task._status = TaskStatus.PAUSED
|
||||
except Exception as e:
|
||||
logging.error(f"task failed, exception occurred: {e}")
|
||||
task._status = TaskStatus.ERROR
|
||||
|
||||
async def TaskManager(taskQueues : Dict[str, list[TaskBase]]):
|
||||
# todo: 处理取消的情况
|
||||
while True:
|
||||
await asyncio.sleep(0.5)
|
||||
for taskQueue in taskQueues.values():
|
||||
notRunningTasks = [task for task in taskQueue if task.worker is None or task.worker.done()]
|
||||
runningTasksNumber = len(taskQueue) - len(notRunningTasks)
|
||||
for task in [task for task in notRunningTasks if task._status == TaskStatus.PENDING]:
|
||||
if runningTasksNumber >= task.maxConcurrentNumber:
|
||||
break
|
||||
task.worker = asyncio.create_task(TaskWorker(task))
|
||||
runningTasksNumber += 1
|
||||
|
||||
|
||||
class FsNode:
|
||||
def __init__(self, id : str, name : str, fatherId : str):
|
||||
self.id = id
|
||||
self.name = name
|
||||
self.fatherId = fatherId
|
||||
self.lastUpdate : datetime = None
|
||||
|
||||
class DirNode(FsNode):
|
||||
def __init__(self, id : str, name : str, fatherId : str):
|
||||
super().__init__(id, name, fatherId)
|
||||
self.childrenId : list[str] = []
|
||||
|
||||
class FileNode(FsNode):
|
||||
def __init__(self, id : str, name : str, fatherId : str):
|
||||
super().__init__(id, name, fatherId)
|
||||
self.url : str = None
|
||||
|
||||
def IsDir(node : FsNode) -> bool:
|
||||
return isinstance(node, DirNode)
|
||||
|
||||
def IsFile(node : FsNode) -> bool:
|
||||
return isinstance(node, FileNode)
|
||||
|
||||
class PikPakToken:
|
||||
def __init__(self, username, password, access_token, refresh_token, user_id):
|
||||
self.username = username
|
||||
self.password = password
|
||||
self.access_token = access_token
|
||||
self.refresh_token = refresh_token
|
||||
self.user_id = user_id
|
||||
|
||||
def to_json(self):
|
||||
return json.dumps(self.__dict__)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_str):
|
||||
data = json.loads(json_str)
|
||||
return cls(**data)
|
||||
|
||||
class PikPakFs:
|
||||
|
||||
async def _pikpak_task_pending(self, task : PikPakTask):
|
||||
pikPakTaskInfo = await self.client.offline_download(task.torrent, task.toDirId)
|
||||
task.remoteTaskId = pikPakTaskInfo["task"]["id"]
|
||||
task.nodeId = pikPakTaskInfo["task"]["file_id"]
|
||||
task.status = PikPakTaskStatus.REMOTE_DOWNLOADING
|
||||
|
||||
async def _pikpak_offline_downloading(self, task : PikPakTask):
|
||||
waitTime = 3
|
||||
while True:
|
||||
await asyncio.sleep(waitTime)
|
||||
status = await self.client.get_task_status(task.remoteTaskId, task.nodeId)
|
||||
if status in {DownloadStatus.not_found, DownloadStatus.not_downloading, DownloadStatus.error}:
|
||||
self.status = PikPakTaskStatus.PENDING
|
||||
raise Exception(f"remote download failed, status: {status}")
|
||||
elif status == DownloadStatus.done:
|
||||
break
|
||||
waitTime = waitTime * 1.5
|
||||
|
||||
fileInfo = await self.client.offline_file_info(file_id=task.nodeId)
|
||||
task.toDirId = fileInfo["parent_id"]
|
||||
task.name = fileInfo["name"]
|
||||
if fileInfo["kind"].endswith("folder"):
|
||||
self.nodes[task.nodeId] = DirNode(task.nodeId, task.name, task.toDirId)
|
||||
else:
|
||||
self.nodes[task.nodeId] = FileNode(task.nodeId, task.name, task.toDirId)
|
||||
|
||||
father = self.GetNodeById(task.toDirId)
|
||||
if father.id is not None and task.nodeId not in father.childrenId:
|
||||
father.childrenId.append(task.nodeId)
|
||||
task.status = PikPakTaskStatus.LOCAL_DOWNLOADING
|
||||
|
||||
async def _pikpak_local_downloading(self, task : PikPakTask):
|
||||
node = self.GetNodeById(task.nodeId)
|
||||
if IsFile(node):
|
||||
fileDownloadTask = FileDownloadTask(task.nodeId, task.id, self.NodeToPath(node, node))
|
||||
fileDownloadTask.handler = self._file_download_task_handler
|
||||
fileDownloadTask.name = task.name
|
||||
self._add_task(fileDownloadTask)
|
||||
elif IsDir(node):
|
||||
# 使用广度优先遍历
|
||||
queue : list[DirNode] = [node]
|
||||
while len(queue) > 0:
|
||||
current = queue.pop(0)
|
||||
await self.Refresh(current)
|
||||
for childId in current.childrenId:
|
||||
child = self.GetNodeById(childId)
|
||||
if IsDir(child):
|
||||
queue.append(child)
|
||||
elif IsFile(child):
|
||||
fileDownloadTask = FileDownloadTask(childId, task.id, self.NodeToPath(child, node))
|
||||
fileDownloadTask.handler = self._file_download_task_handler
|
||||
fileDownloadTask.name = task.name
|
||||
self._add_task(fileDownloadTask)
|
||||
|
||||
# 开始等待下载任务完成
|
||||
while True:
|
||||
fileDownloadTasks = self.taskQueues.get(FileDownloadTask.TAG)
|
||||
myTasks = [myTask for myTask in fileDownloadTasks if myTask.PikPakTaskId == task.id]
|
||||
allNumber = len(myTasks)
|
||||
notCompletedNumber = 0
|
||||
pausedNumber = 0
|
||||
errorNumber = 0
|
||||
for myTask in myTasks:
|
||||
if myTask._status == TaskStatus.PAUSED:
|
||||
pausedNumber += 1
|
||||
if myTask._status == TaskStatus.ERROR:
|
||||
errorNumber += 1
|
||||
if myTask._status in {TaskStatus.PENDING, TaskStatus.RUNNING}:
|
||||
notCompletedNumber += 1
|
||||
|
||||
runningNumber = allNumber - notCompletedNumber - pausedNumber - errorNumber
|
||||
task.progress = f"{runningNumber}/{allNumber} ({pausedNumber}|{errorNumber})"
|
||||
|
||||
if notCompletedNumber > 0:
|
||||
await asyncio.sleep(0.5)
|
||||
continue
|
||||
if errorNumber > 0:
|
||||
raise Exception("file download failed")
|
||||
if pausedNumber > 0:
|
||||
raise asyncio.CancelledError()
|
||||
break
|
||||
|
||||
task.status = PikPakTaskStatus.DONE
|
||||
|
||||
|
||||
async def _pikpak_task_handler(self, task : PikPakTask):
|
||||
while True:
|
||||
if task.status == PikPakTaskStatus.PENDING:
|
||||
await self._pikpak_task_pending(task)
|
||||
elif task.status == PikPakTaskStatus.REMOTE_DOWNLOADING:
|
||||
await self._pikpak_offline_downloading(task)
|
||||
elif task.status == PikPakTaskStatus.LOCAL_DOWNLOADING:
|
||||
await self._pikpak_local_downloading(task)
|
||||
else:
|
||||
break
|
||||
|
||||
async def _file_download_task_handler(self, task : FileDownloadTask):
|
||||
if random.randint(1, 5) == 2:
|
||||
raise asyncio.CancelledError()
|
||||
pass
|
||||
|
||||
def _add_task(self, task : TaskBase):
|
||||
if self.taskQueues.get(task.tag) is None:
|
||||
self.taskQueues[task.tag] = []
|
||||
taskQueue = self.taskQueues[task.tag]
|
||||
for other in taskQueue:
|
||||
if other == task:
|
||||
if other._status != TaskStatus.DONE:
|
||||
other._status = TaskStatus.PENDING
|
||||
return
|
||||
taskQueue.append(task)
|
||||
|
||||
async def GetTaskById(self, taskId : str) -> TaskBase:
|
||||
for taskQueue in self.taskQueues.values():
|
||||
for task in taskQueue:
|
||||
if task.id == taskId:
|
||||
return task
|
||||
return None
|
||||
|
||||
async def StopTask(self, taskId : str):
|
||||
task = await self.GetTaskById(taskId)
|
||||
if task is not None and task._status in {TaskStatus.PENDING, TaskStatus.RUNNING}:
|
||||
if task.worker is not None:
|
||||
task.worker.cancel()
|
||||
task.worker = None
|
||||
|
||||
async def ResumeTask(self, taskId : str):
|
||||
task = await self.GetTaskById(taskId)
|
||||
if task is not None and task._status in {TaskStatus.PAUSED, TaskStatus.ERROR}:
|
||||
task._status = TaskStatus.PENDING
|
||||
|
||||
def Start(self):
|
||||
return asyncio.create_task(TaskManager(self.taskQueues))
|
||||
|
||||
def __init__(self, loginCachePath : str = None, proxy : str = None, rootId = None):
|
||||
self.nodes : Dict[str, FsNode] = {}
|
||||
self.root = DirNode(rootId, "", None)
|
||||
self.currentLocation = self.root
|
||||
|
||||
self.taskQueues : Dict[str, list[TaskBase]] = {}
|
||||
self.loginCachePath = loginCachePath
|
||||
self.proxyConfig = proxy
|
||||
self.client : PikPakApi = None
|
||||
self._try_login_from_cache()
|
||||
|
||||
def _init_client_by_token(self, token : PikPakToken):
|
||||
self._init_client_by_username_and_password(token.username, token.password)
|
||||
self.client.access_token = token.access_token
|
||||
self.client.refresh_token = token.refresh_token
|
||||
self.client.user_id = token.user_id
|
||||
self.client.encode_token()
|
||||
|
||||
def _init_client_by_username_and_password(self, username : str, password : str):
|
||||
httpx_client_args = None
|
||||
if self.proxyConfig != None:
|
||||
httpx_client_args = {
|
||||
"proxy": self.proxyConfig,
|
||||
"transport": httpx.AsyncHTTPTransport()
|
||||
}
|
||||
|
||||
self.client = PikPakApi(
|
||||
username = username,
|
||||
password = password,
|
||||
httpx_client_args=httpx_client_args)
|
||||
|
||||
def _try_login_from_cache(self):
|
||||
if self.loginCachePath is None:
|
||||
return
|
||||
if not os.path.exists(self.loginCachePath):
|
||||
return
|
||||
with open(self.loginCachePath, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
token = PikPakToken.from_json(content)
|
||||
self._init_client_by_token(token)
|
||||
logging.info("successfully load login info from cache")
|
||||
|
||||
def _dump_login_info(self):
|
||||
if self.loginCachePath is None:
|
||||
return
|
||||
with open(self.loginCachePath, 'w', encoding='utf-8') as file:
|
||||
token = PikPakToken(self.client.username, self.client.password, self.client.access_token, self.client.refresh_token, self.client.user_id)
|
||||
file.write(token.to_json())
|
||||
logging.info("successfully dump login info to cache")
|
||||
|
||||
def _is_ancestors_of(self, nodeA : FsNode, nodeB : FsNode) -> bool:
|
||||
if nodeB is nodeA:
|
||||
return False
|
||||
if nodeA is self.root:
|
||||
return True
|
||||
while nodeB.fatherId != self.root.id:
|
||||
nodeB = self.nodes[nodeB.fatherId]
|
||||
if nodeB is nodeA:
|
||||
return True
|
||||
return False
|
||||
|
||||
def GetNodeById(self, id : str) -> FsNode:
|
||||
if id == self.root.id:
|
||||
return self.root
|
||||
if id not in self.nodes:
|
||||
return None
|
||||
return self.nodes[id]
|
||||
|
||||
def GetFatherNode(self, node : FsNode) -> FsNode:
|
||||
if node is self.root:
|
||||
return self.root
|
||||
return self.GetNodeById(node.fatherId)
|
||||
|
||||
def FindChildInDirByName(self, dir : DirNode, name : str):
|
||||
if dir is self.root and name == "":
|
||||
return self.root
|
||||
for childId in dir.childrenId:
|
||||
node = self.nodes[childId]
|
||||
if name == node.name:
|
||||
return node
|
||||
return None
|
||||
|
||||
async def Refresh(self, node : FsNode):
|
||||
if IsDir(node):
|
||||
if node.lastUpdate != None:
|
||||
return
|
||||
next_page_token = None
|
||||
childrenInfo = []
|
||||
while True:
|
||||
dirInfo = await self.client.file_list(parent_id = node.id, next_page_token=next_page_token)
|
||||
next_page_token = dirInfo["next_page_token"]
|
||||
currentPageNodes = dirInfo["files"]
|
||||
childrenInfo.extend(currentPageNodes)
|
||||
if next_page_token is None or next_page_token == "":
|
||||
break
|
||||
node.childrenId.clear()
|
||||
|
||||
for childInfo in childrenInfo:
|
||||
child : FsNode = None
|
||||
id = childInfo["id"]
|
||||
name = childInfo["name"]
|
||||
fatherId = node.id
|
||||
if id in self.nodes:
|
||||
child = self.nodes[id]
|
||||
else:
|
||||
child = DirNode(id, name, fatherId) if childInfo["kind"].endswith("folder") else FileNode(id, name, fatherId)
|
||||
self.nodes[id] = child
|
||||
child.name = name
|
||||
child.fatherId = fatherId
|
||||
node.childrenId.append(id)
|
||||
elif IsFile(node):
|
||||
result = await self.client.get_download_url(node.id)
|
||||
node.url = result["web_content_link"]
|
||||
|
||||
node.lastUpdate = datetime.now()
|
||||
|
||||
async def PathToNode(self, path : str) -> FsNode:
|
||||
father, sonName = await self.PathToFatherNodeAndNodeName(path)
|
||||
if sonName == "":
|
||||
return father
|
||||
if not IsDir(father):
|
||||
return None
|
||||
return self.FindChildInDirByName(father, sonName)
|
||||
|
||||
async def PathToFatherNodeAndNodeName(self, path : str) -> tuple[FsNode, str]:
|
||||
pathWalker = PathWalker(path)
|
||||
father : FsNode = None
|
||||
sonName : str = None
|
||||
current = self.root if pathWalker.IsAbsolute() else self.currentLocation
|
||||
|
||||
for spot in pathWalker.Walk():
|
||||
if current is None:
|
||||
father = None
|
||||
break
|
||||
if spot == "..":
|
||||
current = self.GetFatherNode(current)
|
||||
continue
|
||||
father = current
|
||||
if not IsDir(current):
|
||||
current = None
|
||||
continue
|
||||
await self.Refresh(current)
|
||||
if spot == ".":
|
||||
continue
|
||||
sonName = spot
|
||||
current = self.FindChildInDirByName(current, spot)
|
||||
|
||||
if current != None:
|
||||
father = self.GetFatherNode(current)
|
||||
sonName = current.name
|
||||
|
||||
return father, sonName
|
||||
|
||||
def NodeToPath(self, node : FsNode, root : FsNode = None) -> str:
|
||||
if root is None:
|
||||
root = self.root
|
||||
if node is root:
|
||||
return "/"
|
||||
spots : list[str] = []
|
||||
current = node
|
||||
while current is not root:
|
||||
spots.append(current.name)
|
||||
current = self.GetFatherNode(current)
|
||||
spots.append("")
|
||||
return "/".join(reversed(spots))
|
||||
|
||||
# commands #
|
||||
async def Login(self, username : str = None, password : str = None) -> None:
|
||||
if self.client != None and username is None and password is None:
|
||||
username = self.client.username
|
||||
password = self.client.password
|
||||
|
||||
if username == None and password == None:
|
||||
raise Exception("Username and password are required")
|
||||
|
||||
self._init_client_by_username_and_password(username, password)
|
||||
await self.client.login()
|
||||
self._dump_login_info()
|
||||
|
||||
async def MakeDir(self, node : DirNode, name : str) -> DirNode:
|
||||
result = await self.client.create_folder(name, node.id)
|
||||
id = result["file"]["id"]
|
||||
name = result["file"]["name"]
|
||||
newDir = DirNode(id, name, node.id)
|
||||
self.nodes[id] = newDir
|
||||
node.childrenId.append(id)
|
||||
return newDir
|
||||
|
||||
async def Download(self, url : str, dirNode : DirNode) -> PikPakTask :
|
||||
task = PikPakTask(url, dirNode.id)
|
||||
task.handler = self._pikpak_task_handler
|
||||
self._add_task(task)
|
||||
return task
|
||||
|
||||
async def Pull(self, node : FsNode) -> PikPakTask:
|
||||
task = PikPakTask("", node.fatherId, node.id, PikPakTaskStatus.LOCAL_DOWNLOADING)
|
||||
task.name = node.name
|
||||
task.handler = self._pikpak_task_handler
|
||||
self._add_task(task)
|
||||
return task
|
||||
|
||||
async def QueryPikPakTasks(self, filterStatus : TaskStatus = None) -> list[PikPakTask]:
|
||||
if PikPakTask.TAG not in self.taskQueues:
|
||||
return []
|
||||
taskQueue = self.taskQueues[PikPakTask.TAG]
|
||||
if filterStatus is None:
|
||||
return taskQueue
|
||||
return [task for task in taskQueue if task._status == filterStatus]
|
||||
|
||||
async def QueryFileDownloadTasks(self, filterStatus : TaskStatus = None) -> list[FileDownloadTask]:
|
||||
if FileDownloadTask.TAG not in self.taskQueues:
|
||||
return []
|
||||
taskQueue = self.taskQueues[FileDownloadTask.TAG]
|
||||
if filterStatus is None:
|
||||
return taskQueue
|
||||
return [task for task in taskQueue if task._status == filterStatus]
|
||||
|
||||
async def Delete(self, nodes : list[FsNode]) -> None:
|
||||
nodeIds = [node.id for node in nodes]
|
||||
await self.client.delete_to_trash(nodeIds)
|
||||
for node in nodes:
|
||||
father = self.GetFatherNode(node)
|
||||
father.childrenId.remove(node.id)
|
||||
295
TaskManager.py
Normal file
295
TaskManager.py
Normal file
@ -0,0 +1,295 @@
|
||||
from enum import Enum
|
||||
from typing import Awaitable, Callable, Dict
|
||||
import asyncio
|
||||
import logging
|
||||
import shortuuid
|
||||
from PikPakFileSystem import PikPakFileSystem
|
||||
from pikpakapi import DownloadStatus
|
||||
import random
|
||||
|
||||
class TaskStatus(Enum):
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
DONE = "done"
|
||||
ERROR = "error"
|
||||
PAUSED = "paused"
|
||||
|
||||
class TorrentTaskStatus(Enum):
|
||||
PENDING = "pending"
|
||||
REMOTE_DOWNLOADING = "remote"
|
||||
LOCAL_DOWNLOADING = "local"
|
||||
DONE = "done"
|
||||
|
||||
class FileDownloadTaskStatus(Enum):
|
||||
PENDING = "pending"
|
||||
DOWNLOADING = "downloading"
|
||||
DONE = "done"
|
||||
|
||||
class TaskBase:
|
||||
TAG = ""
|
||||
MAX_CONCURRENT_NUMBER = 5
|
||||
|
||||
def __init__(self, client : PikPakFileSystem):
|
||||
self.id : str = shortuuid.uuid()
|
||||
self.status : TaskStatus = TaskStatus.PENDING
|
||||
self.worker : asyncio.Task = None
|
||||
self.handler : Callable[..., Awaitable] = None
|
||||
self.client : PikPakFileSystem = client
|
||||
|
||||
def Resume(self):
|
||||
if self.status in {TaskStatus.PAUSED, TaskStatus.ERROR}:
|
||||
self.status = TaskStatus.PENDING
|
||||
|
||||
|
||||
class TorrentTask(TaskBase):
|
||||
TAG = "TorrentTask"
|
||||
MAX_CONCURRENT_NUMBER = 5
|
||||
|
||||
def __init__(self, torrent : str):
|
||||
super().__init__(self)
|
||||
self.torrent_status : TorrentTaskStatus = TorrentTaskStatus.PENDING
|
||||
self.torrent : str = torrent
|
||||
self.remote_full_path : str = None
|
||||
self.remote_base_path : str = None
|
||||
self.info : str = ""
|
||||
|
||||
# 和PikPak交互需要的信息
|
||||
self.node_id : str = None
|
||||
self.task_id : str = None
|
||||
|
||||
|
||||
class FileDownloadTask(TaskBase):
|
||||
TAG = "FileDownloadTask"
|
||||
MAX_CONCURRENT_NUMBER = 5
|
||||
|
||||
def __init__(self, remote_path : str, owner_id : str):
|
||||
super().__init__(self)
|
||||
self.file_download_status : FileDownloadTaskStatus = FileDownloadTaskStatus.PENDING
|
||||
self.remote_path : str = remote_path
|
||||
self.owner_id : str = owner_id
|
||||
|
||||
async def TaskWorker(task : TaskBase):
|
||||
try:
|
||||
if task.status != TaskStatus.PENDING:
|
||||
return
|
||||
task.status = TaskStatus.RUNNING
|
||||
await task.handler(task)
|
||||
task.status = TaskStatus.DONE
|
||||
except asyncio.CancelledError:
|
||||
task.status = TaskStatus.PAUSED
|
||||
except Exception as e:
|
||||
logging.error(f"task failed, exception occurred: {e}")
|
||||
task.status = TaskStatus.ERROR
|
||||
|
||||
class TaskManager:
|
||||
#region 内部实现
|
||||
def __init__(self, client : PikPakFileSystem):
|
||||
self.taskQueues : Dict[str, list[TaskBase]] = {}
|
||||
self.loop : asyncio.Task = None
|
||||
self.client = client
|
||||
|
||||
async def _loop(self):
|
||||
while True:
|
||||
await asyncio.sleep(0.5)
|
||||
for taskQueue in self.taskQueues.values():
|
||||
notRunningTasks = [task for task in taskQueue if task.worker is None or task.worker.done()]
|
||||
runningTasksNumber = len(taskQueue) - len(notRunningTasks)
|
||||
for task in [task for task in notRunningTasks if task.status == TaskStatus.PENDING]:
|
||||
if runningTasksNumber >= task.MAX_CONCURRENT_NUMBER:
|
||||
break
|
||||
task.worker = asyncio.create_task(TaskWorker(task))
|
||||
runningTasksNumber += 1
|
||||
|
||||
async def _get_task_by_id(self, task_id : str) -> TaskBase:
|
||||
for queue in self.taskQueues.values():
|
||||
for task in queue:
|
||||
if task.id == task_id:
|
||||
return task
|
||||
return None
|
||||
|
||||
#region 远程下载部分
|
||||
|
||||
async def _append_task(self, task : TaskBase):
|
||||
queue = self.taskQueues.get(task.TAG, [])
|
||||
queue.append(task)
|
||||
self.taskQueues[task.TAG] = queue
|
||||
|
||||
async def _get_torrent_queue(self):
|
||||
if TorrentTask.TAG not in self.taskQueues:
|
||||
self.taskQueues[TorrentTask.TAG] = []
|
||||
return self.taskQueues[TorrentTask.TAG]
|
||||
|
||||
async def _get_file_download_queue(self, owner_id : str):
|
||||
if FileDownloadTask.TAG not in self.taskQueues:
|
||||
self.taskQueues[FileDownloadTask.TAG] = []
|
||||
queue = self.taskQueues[FileDownloadTask.TAG]
|
||||
return [task for task in queue if task.owner_id == owner_id]
|
||||
|
||||
async def _on_torrent_task_pending(self, task : TorrentTask):
|
||||
task.node_id, task.task_id = await self.client.RemoteDownload(task.torrent, task.remote_base_path)
|
||||
task.torrent_status = TorrentTaskStatus.REMOTE_DOWNLOADING
|
||||
|
||||
async def _on_torrent_task_offline_downloading(self, task : TorrentTask):
|
||||
wait_seconds = 3
|
||||
while True:
|
||||
status = await self.client.QueryTaskStatus(task.task_id, task.node_id)
|
||||
if status in {DownloadStatus.not_found, DownloadStatus.not_downloading, DownloadStatus.error}:
|
||||
task.torrent_status = TorrentTaskStatus.PENDING
|
||||
raise Exception(f"remote download failed, status: {status}")
|
||||
elif status == DownloadStatus.done:
|
||||
break
|
||||
await asyncio.sleep(wait_seconds)
|
||||
wait_seconds = wait_seconds * 1.5
|
||||
|
||||
task.remote_full_path = await self.client.UpdateDirectory(task.remote_base_path, task.node_id)
|
||||
task.torrent_status = TorrentTaskStatus.LOCAL_DOWNLOADING
|
||||
|
||||
async def _on_torrent_local_downloading(self, task : TorrentTask):
|
||||
path = task.remote_full_path
|
||||
|
||||
if not await self.client.IsDir(path):
|
||||
await self._init_file_download_task(path, task.id)
|
||||
else:
|
||||
# 使用广度优先遍历
|
||||
queue : list[str] = [path]
|
||||
while len(queue) > 0:
|
||||
current_path = queue.pop(0)
|
||||
for child_name in await self.client.GetChildrenNames(current_path, False):
|
||||
child_path = await self.client.JoinPath(current_path, child_name)
|
||||
if await self.client.IsDir(child_path):
|
||||
queue.append(child_path)
|
||||
else:
|
||||
await self._init_file_download_task(child_path, task.id)
|
||||
|
||||
# 开始等待下载任务完成
|
||||
while True:
|
||||
file_download_tasks = await self._get_file_download_queue(task.id)
|
||||
all_number = len(file_download_tasks)
|
||||
not_completed_number = 0
|
||||
paused_number = 0
|
||||
error_number = 0
|
||||
for file_download_task in file_download_tasks:
|
||||
if file_download_task.status == TaskStatus.PAUSED:
|
||||
paused_number += 1
|
||||
if file_download_task.status == TaskStatus.ERROR:
|
||||
error_number += 1
|
||||
if file_download_task.status in {TaskStatus.PENDING, TaskStatus.RUNNING}:
|
||||
not_completed_number += 1
|
||||
|
||||
running_number = all_number - not_completed_number - paused_number - error_number
|
||||
task.info = f"{running_number}/{all_number} ({paused_number}|{error_number})"
|
||||
|
||||
if not_completed_number > 0:
|
||||
await asyncio.sleep(0.5)
|
||||
continue
|
||||
if error_number > 0:
|
||||
raise Exception("file download failed")
|
||||
if paused_number > 0:
|
||||
raise asyncio.CancelledError()
|
||||
break
|
||||
|
||||
task.torrent_status = TorrentTaskStatus.DONE
|
||||
|
||||
async def _on_torrent_task_cancelled(self, task : TorrentTask):
|
||||
file_download_tasks = await self._get_file_download_queue(task.id)
|
||||
for file_download_task in file_download_tasks:
|
||||
if file_download_task.worker is not None:
|
||||
file_download_task.worker.cancel()
|
||||
|
||||
async def _torrent_task_handler(self, task : TorrentTask):
|
||||
try:
|
||||
while True:
|
||||
if task.torrent_status == TorrentTaskStatus.PENDING:
|
||||
await self._on_torrent_task_pending(task)
|
||||
elif task.torrent_status == TorrentTaskStatus.REMOTE_DOWNLOADING:
|
||||
await self._on_torrent_task_offline_downloading(task)
|
||||
elif task.torrent_status == TorrentTaskStatus.LOCAL_DOWNLOADING:
|
||||
await self._on_torrent_local_downloading(task)
|
||||
else:
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
await self._on_torrent_task_cancelled(task)
|
||||
raise
|
||||
#endregion
|
||||
|
||||
|
||||
#region 文件下载部分
|
||||
async def _init_file_download_task(self, remote_path : str, owner_id : str) -> str:
|
||||
queue = await self._get_file_download_queue(owner_id)
|
||||
for task in queue:
|
||||
if not isinstance(task, FileDownloadTask):
|
||||
continue
|
||||
if task.remote_path == remote_path:
|
||||
if task.status in {TaskStatus.PAUSED, TaskStatus.ERROR}:
|
||||
task.status = TaskStatus.PENDING
|
||||
return task.id
|
||||
task = FileDownloadTask(remote_path, owner_id)
|
||||
task.handler = self._file_download_task_handler
|
||||
await self._append_task(task)
|
||||
return task.id
|
||||
|
||||
async def _file_download_task_handler(self, task : FileDownloadTask):
|
||||
await asyncio.sleep(30)
|
||||
if random.randint(1, 5) == 2:
|
||||
raise asyncio.CancelledError()
|
||||
if random.randint(1, 5) == 3:
|
||||
raise Exception("random error")
|
||||
pass
|
||||
|
||||
#endregion
|
||||
|
||||
#endregion
|
||||
|
||||
#region 对外接口
|
||||
|
||||
def Start(self):
|
||||
# todo: 从文件中恢复任务
|
||||
if self.loop is None:
|
||||
self.loop = asyncio.create_task(self._loop())
|
||||
|
||||
def Stop(self):
|
||||
if self.loop is not None:
|
||||
self.loop.cancel()
|
||||
self.loop = None
|
||||
# todo: 保存任务到文件
|
||||
|
||||
async def CreateTorrentTask(self, torrent : str, remote_base_path : str) -> str:
|
||||
task = TorrentTask(torrent)
|
||||
task.remote_base_path = remote_base_path
|
||||
task.handler = self._torrent_task_handler
|
||||
await self._append_task(task)
|
||||
return task.id
|
||||
|
||||
async def PullRemote(self, remote_full_path : str) -> str:
|
||||
if not await self.client.IfExists(remote_full_path):
|
||||
raise Exception("target not found")
|
||||
queue = await self._get_torrent_queue()
|
||||
for task in queue:
|
||||
if not isinstance(task, TorrentTask):
|
||||
continue
|
||||
if task.remote_full_path == remote_full_path:
|
||||
return task.id
|
||||
task = TorrentTask(None)
|
||||
task.remote_full_path = remote_full_path
|
||||
task.handler = self._torrent_task_handler
|
||||
task.torrent_status = TorrentTaskStatus.LOCAL_DOWNLOADING
|
||||
await self._append_task(task)
|
||||
return task.id
|
||||
|
||||
async def QueryTasks(self, tag : str, filter_status : TaskStatus = None):
|
||||
queue = self.taskQueues.get(tag, [])
|
||||
if filter_status is None:
|
||||
return queue
|
||||
return [task for task in queue if task.status == filter_status]
|
||||
|
||||
async def StopTask(self, task_id : str):
|
||||
task = await self._get_task_by_id(task_id)
|
||||
if task is not None and task.worker is not None:
|
||||
task.worker.cancel()
|
||||
|
||||
async def ResumeTask(self, task_id : str):
|
||||
task = await self._get_task_by_id(task_id)
|
||||
if task is not None:
|
||||
task.Resume()
|
||||
|
||||
#endregion
|
||||
162
main.py
162
main.py
@ -4,11 +4,11 @@ from functools import wraps
|
||||
import logging
|
||||
import threading
|
||||
import colorlog
|
||||
from PikPakFs import PikPakFs, IsDir, IsFile, TaskStatus
|
||||
from PikPakFileSystem import PikPakFileSystem
|
||||
import os
|
||||
from tabulate import tabulate
|
||||
import wcwidth
|
||||
import types
|
||||
from TaskManager import TaskManager, TaskStatus, TorrentTask, FileDownloadTask
|
||||
|
||||
LogFormatter = colorlog.ColoredFormatter(
|
||||
"%(log_color)s%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
@ -34,7 +34,7 @@ def setup_logging():
|
||||
|
||||
setup_logging()
|
||||
MainLoop : asyncio.AbstractEventLoop = None
|
||||
Client = PikPakFs("token.json", proxy="http://127.0.0.1:7897")
|
||||
Client = PikPakFileSystem(auth_cache_path = "token.json", proxy_address="http://127.0.0.1:7897")
|
||||
|
||||
class RunSync:
|
||||
_current_task : asyncio.Task = None
|
||||
@ -69,18 +69,19 @@ class RunSync:
|
||||
else:
|
||||
return types.MethodType(self, instance)
|
||||
|
||||
class Console(cmd2.Cmd):
|
||||
class App(cmd2.Cmd):
|
||||
#region Console设置
|
||||
def _io_worker(self, loop):
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_forever()
|
||||
|
||||
async def Input(self, prompt):
|
||||
async def input(self, prompt):
|
||||
async def _input(prompt):
|
||||
return self._read_command_line(prompt)
|
||||
future = asyncio.run_coroutine_threadsafe(_input(prompt), self.ioLoop)
|
||||
return await asyncio.wrap_future(future)
|
||||
|
||||
async def Print(self, *args, **kwargs):
|
||||
async def print(self, *args, **kwargs):
|
||||
async def _print(*args, **kwargs):
|
||||
print(*args, **kwargs)
|
||||
future = asyncio.run_coroutine_threadsafe(_print(*args, **kwargs), self.ioLoop)
|
||||
@ -93,6 +94,8 @@ class Console(cmd2.Cmd):
|
||||
self.log_handler.setLevel(logging.CRITICAL)
|
||||
logging.getLogger().addHandler(self.log_handler)
|
||||
|
||||
self.task_manager = TaskManager(Client)
|
||||
|
||||
def preloop(self):
|
||||
# 1. 设置忽略SIGINT
|
||||
import signal
|
||||
@ -109,6 +112,9 @@ class Console(cmd2.Cmd):
|
||||
self.saved_readline_settings = None
|
||||
with self.sigint_protection:
|
||||
self.saved_readline_settings = self._set_up_cmd2_readline()
|
||||
|
||||
# 4. 启动任务管理器
|
||||
self.task_manager.Start()
|
||||
|
||||
def postloop(self):
|
||||
# 1. 还原console设置
|
||||
@ -120,8 +126,13 @@ class Console(cmd2.Cmd):
|
||||
# https://stackoverflow.com/questions/51642267/asyncio-how-do-you-use-run-forever
|
||||
self.ioLoop.call_soon_threadsafe(self.ioLoop.stop)
|
||||
self.ioThread.join()
|
||||
|
||||
# commands #
|
||||
|
||||
# 3. 停止任务管理器
|
||||
self.task_manager.Stop()
|
||||
|
||||
#endregion
|
||||
|
||||
#region 所有命令
|
||||
def do_logging_off(self, args):
|
||||
"""
|
||||
Disable logging
|
||||
@ -153,27 +164,21 @@ class Console(cmd2.Cmd):
|
||||
Login to pikpak
|
||||
"""
|
||||
await Client.Login(args.username, args.password)
|
||||
await self.Print("Logged in successfully")
|
||||
await self.print("Logged in successfully")
|
||||
|
||||
async def _path_completer(self, text, line, begidx, endidx, ignoreFiles):
|
||||
father, sonName = await Client.PathToFatherNodeAndNodeName(text)
|
||||
if not IsDir(father):
|
||||
return []
|
||||
father_path, son_name = await Client.SplitPath(text)
|
||||
children_names = await Client.GetChildrenNames(father_path, ignoreFiles)
|
||||
matches = []
|
||||
matchesNode = []
|
||||
for childId in father.childrenId:
|
||||
child = Client.GetNodeById(childId)
|
||||
if ignoreFiles and IsFile(child):
|
||||
continue
|
||||
if child.name.startswith(sonName):
|
||||
self.display_matches.append(child.name)
|
||||
if sonName == "":
|
||||
matches.append(text + child.name)
|
||||
elif text.endswith(sonName):
|
||||
matches.append(text[:text.rfind(sonName)] + child.name)
|
||||
matchesNode.append(child)
|
||||
if len(matchesNode) == 1 and IsDir(matchesNode[0]):
|
||||
if matches[0] == sonName:
|
||||
for child_name in children_names:
|
||||
if child_name.startswith(son_name):
|
||||
self.display_matches.append(child_name)
|
||||
if son_name == "":
|
||||
matches.append(text + child_name)
|
||||
elif text.endswith(son_name):
|
||||
matches.append(text[:text.rfind(son_name)] + child_name)
|
||||
if len(matches) == 1 and await Client.IsDir(father_path + matches[0]):
|
||||
if matches[0] == son_name:
|
||||
matches[0] += "/"
|
||||
self.allow_appended_space = False
|
||||
self.allow_closing_quote = False
|
||||
@ -184,49 +189,37 @@ class Console(cmd2.Cmd):
|
||||
return await self._path_completer(text, line, begidx, endidx, False)
|
||||
|
||||
ls_parser = cmd2.Cmd2ArgumentParser()
|
||||
ls_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
|
||||
ls_parser.add_argument("path", help="path", default="", nargs="?")
|
||||
@cmd2.with_argparser(ls_parser)
|
||||
@RunSync
|
||||
async def do_ls(self, args):
|
||||
"""
|
||||
List files in a directory
|
||||
"""
|
||||
node = args.path
|
||||
if node is None:
|
||||
await self.Print("Invalid path")
|
||||
return
|
||||
await Client.Refresh(node)
|
||||
if IsDir(node):
|
||||
for childId in node.childrenId:
|
||||
child = Client.GetNodeById(childId)
|
||||
await self.Print(child.name)
|
||||
elif IsFile(node):
|
||||
await self.Print(f"{node.name}: {node.url}")
|
||||
if await Client.IsDir(args.path):
|
||||
for child_name in await Client.GetChildrenNames(args.path, False):
|
||||
await self.print(child_name)
|
||||
|
||||
@RunSync
|
||||
async def complete_cd(self, text, line, begidx, endidx):
|
||||
return await self._path_completer(text, line, begidx, endidx, True)
|
||||
|
||||
cd_parser = cmd2.Cmd2ArgumentParser()
|
||||
cd_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
|
||||
cd_parser.add_argument("path", help="path", default="", nargs="?")
|
||||
@cmd2.with_argparser(cd_parser)
|
||||
@RunSync
|
||||
async def do_cd(self, args):
|
||||
"""
|
||||
Change directory
|
||||
"""
|
||||
node = args.path
|
||||
if not IsDir(node):
|
||||
await self.Print("Invalid directory")
|
||||
return
|
||||
Client.currentLocation = node
|
||||
await Client.SetCwd(args.path)
|
||||
|
||||
@RunSync
|
||||
async def do_cwd(self, args):
|
||||
"""
|
||||
Print current working directory
|
||||
"""
|
||||
await self.Print(Client.NodeToPath(Client.currentLocation))
|
||||
await self.print(await Client.GetCwd())
|
||||
|
||||
def do_clear(self, args):
|
||||
"""
|
||||
@ -239,7 +232,7 @@ class Console(cmd2.Cmd):
|
||||
return await self._path_completer(text, line, begidx, endidx, False)
|
||||
|
||||
rm_parser = cmd2.Cmd2ArgumentParser()
|
||||
rm_parser.add_argument("paths", help="paths", default="", nargs="+", type=RunSync(Client.PathToNode))
|
||||
rm_parser.add_argument("paths", help="paths", default="", nargs="+")
|
||||
@cmd2.with_argparser(rm_parser)
|
||||
@RunSync
|
||||
async def do_rm(self, args):
|
||||
@ -253,56 +246,44 @@ class Console(cmd2.Cmd):
|
||||
return await self._path_completer(text, line, begidx, endidx, True)
|
||||
|
||||
mkdir_parser = cmd2.Cmd2ArgumentParser()
|
||||
mkdir_parser.add_argument("path_and_son", help="path and son", default="", nargs="?", type=RunSync(Client.PathToFatherNodeAndNodeName))
|
||||
mkdir_parser.add_argument("path", help="new directory path")
|
||||
@cmd2.with_argparser(mkdir_parser)
|
||||
@RunSync
|
||||
async def do_mkdir(self, args):
|
||||
"""
|
||||
Create a directory
|
||||
"""
|
||||
father, sonName = args.path_and_son
|
||||
if not IsDir(father) or sonName == "" or sonName == None:
|
||||
await self.Print("Invalid path")
|
||||
return
|
||||
child = Client.FindChildInDirByName(father, sonName)
|
||||
if child is not None:
|
||||
await self.Print("Path already exists")
|
||||
return
|
||||
await Client.MakeDir(father, sonName)
|
||||
await Client.MakeDir(args.path)
|
||||
|
||||
download_parser = cmd2.Cmd2ArgumentParser()
|
||||
download_parser.add_argument("url", help="url")
|
||||
download_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
|
||||
download_parser.add_argument("torrent", help="torrent")
|
||||
@cmd2.with_argparser(download_parser)
|
||||
@RunSync
|
||||
async def do_download(self, args):
|
||||
"""
|
||||
Download a file or directory
|
||||
"""
|
||||
node = args.path
|
||||
if not IsDir(node):
|
||||
await self.Print("Invalid directory")
|
||||
return
|
||||
task = await Client.Download(args.url, node)
|
||||
await self.Print(f"Task {task.id} created")
|
||||
task_id = await self.task_manager.CreateTorrentTask(args.torrent, await Client.GetCwd())
|
||||
await self.print(f"Task {task_id} created")
|
||||
|
||||
@RunSync
|
||||
async def complete_pull(self, text, line, begidx, endidx):
|
||||
return await self._path_completer(text, line, begidx, endidx, False)
|
||||
|
||||
pull_parser = cmd2.Cmd2ArgumentParser()
|
||||
pull_parser.add_argument("target", help="pull target", type=RunSync(Client.PathToNode))
|
||||
pull_parser.add_argument("target", help="pull target")
|
||||
@cmd2.with_argparser(pull_parser)
|
||||
@RunSync
|
||||
async def do_pull(self, args):
|
||||
"""
|
||||
Pull a file or directory
|
||||
"""
|
||||
await Client.Pull(args.target)
|
||||
task_id = await self.task_manager.PullRemote(await Client.ToFullPath(args.target))
|
||||
await self.print(f"Task {task_id} created")
|
||||
|
||||
|
||||
query_parser = cmd2.Cmd2ArgumentParser()
|
||||
query_parser.add_argument("-t", "--type", help="type", nargs="?", choices=["pikpak", "filedownload"], default="pikpak")
|
||||
query_parser.add_argument("-t", "--type", help="type", nargs="?", choices=["torrent", "file"], default="torrent")
|
||||
query_parser.add_argument("-f", "--filter", help="filter", nargs="?", choices=[member.value for member in TaskStatus])
|
||||
@cmd2.with_argparser(query_parser)
|
||||
@RunSync
|
||||
@ -310,21 +291,21 @@ class Console(cmd2.Cmd):
|
||||
"""
|
||||
Query All Tasks
|
||||
"""
|
||||
if args.type == "pikpak":
|
||||
tasks = await Client.QueryPikPakTasks(TaskStatus(args.filter) if args.filter is not None else None)
|
||||
filter_status = TaskStatus(args.filter) if args.filter is not None else None
|
||||
if args.type == "torrent":
|
||||
tasks = await self.task_manager.QueryTasks(TorrentTask.TAG, filter_status)
|
||||
# 格式化输出所有task信息id，status，lastStatus的信息，输出表格
|
||||
table = [[task.id, task._status.value, task.status.value, task.progress] for task in tasks]
|
||||
table = [[task.id, task.status.value, task.torrent_status.value, task.info] for task in tasks if isinstance(task, TorrentTask)]
|
||||
headers = ["id", "status", "details", "progress"]
|
||||
await self.Print(tabulate(table, headers, tablefmt="grid"))
|
||||
elif args.type == "filedownload":
|
||||
tasks = await Client.QueryFileDownloadTasks(TaskStatus(args.filter) if args.filter is not None else None)
|
||||
# 格式化输出所有task信息id，status，lastStatus的信息，输出表格
|
||||
table = [[task.id, task._status.value, task.status.value, task.relativePath] for task in tasks]
|
||||
headers = ["id", "status", "details", "path"]
|
||||
await self.Print(tabulate(table, headers, tablefmt="grid"))
|
||||
await self.print(tabulate(table, headers, tablefmt="grid"))
|
||||
elif args.type == "file":
|
||||
tasks = await self.task_manager.QueryTasks(FileDownloadTask.TAG, filter_status)
|
||||
table = [[task.id, task.status.value, task.file_download_status, task.remote_path] for task in tasks if isinstance(task, FileDownloadTask)]
|
||||
headers = ["id", "status", "details", "remote_path"]
|
||||
await self.print(tabulate(table, headers, tablefmt="grid"))
|
||||
|
||||
taskid_parser = cmd2.Cmd2ArgumentParser()
|
||||
taskid_parser.add_argument("taskId", help="taskId")
|
||||
taskid_parser.add_argument("task_id", help="task id")
|
||||
|
||||
@cmd2.with_argparser(taskid_parser)
|
||||
@RunSync
|
||||
@ -332,7 +313,7 @@ class Console(cmd2.Cmd):
|
||||
"""
|
||||
Stop a task
|
||||
"""
|
||||
await Client.StopTask(args.taskId)
|
||||
await self.task_manager.StopTask(args.task_id)
|
||||
|
||||
@cmd2.with_argparser(taskid_parser)
|
||||
@RunSync
|
||||
@ -340,27 +321,30 @@ class Console(cmd2.Cmd):
|
||||
"""
|
||||
Resume a task
|
||||
"""
|
||||
await Client.ResumeTask(args.taskId)
|
||||
await self.task_manager.ResumeTask(args.task_id)
|
||||
|
||||
#endregion
|
||||
|
||||
|
||||
#region APP入口
|
||||
async def mainLoop():
|
||||
global MainLoop, Client
|
||||
global MainLoop
|
||||
MainLoop = asyncio.get_running_loop()
|
||||
clientWorker = Client.Start()
|
||||
app = App()
|
||||
|
||||
console = Console()
|
||||
console.preloop()
|
||||
app.preloop()
|
||||
try:
|
||||
stop = False
|
||||
while not stop:
|
||||
line = await console.Input(console.prompt)
|
||||
line = await app.input(app.prompt)
|
||||
try:
|
||||
stop = console.onecmd_plus_hooks(line)
|
||||
stop = app.onecmd_plus_hooks(line)
|
||||
except asyncio.CancelledError:
|
||||
await console.Print("^C: Task cancelled")
|
||||
await app.print("^C: Task cancelled")
|
||||
finally:
|
||||
console.postloop()
|
||||
clientWorker.cancel()
|
||||
app.postloop()
|
||||
|
||||
if __name__ == "__main__":
|
||||
nest_asyncio.apply()
|
||||
asyncio.run(mainLoop())
|
||||
#endregion
|
||||
16
utils.py
16
utils.py
@ -1,16 +0,0 @@
|
||||
class PathWalker():
|
||||
def __init__(self, path : str, sep : str = "/"):
|
||||
self._path_spots : list[str] = []
|
||||
if not path.startswith(sep):
|
||||
self._path_spots.append(".")
|
||||
path_spots : list[str] = path.split(sep)
|
||||
self._path_spots.extend(path_spots)
|
||||
|
||||
def IsAbsolute(self) -> bool:
|
||||
return len(self._path_spots) == 0 or self._path_spots[0] != "."
|
||||
|
||||
def AppendSpot(self, spot):
|
||||
self._path_spots.append(spot)
|
||||
|
||||
def Walk(self) -> list[str]:
|
||||
return self._path_spots
|
||||
Loading…
x
Reference in New Issue
Block a user