changes

gramup/main.py

@@ -1,17 +1,18 @@
 import asyncio
 import glob
 import os
+from pathlib import Path
 from typing import List
 
-from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi import FastAPI, Request, BackgroundTasks, HTTPException
 from starlette.middleware.cors import CORSMiddleware
 from starlette.staticfiles import StaticFiles
 from starlette.websockets import WebSocket
 from telethon import TelegramClient
 
-from constants import API_ID, API_HASH
+from constants import API_ID, API_HASH, DB_PATH
 from gramup.models import Task
-from gramup.tasks import pull_all_to_db
+from gramup.tasks import pull_all_to_db, upload
 
 client = TelegramClient('anon', API_ID, API_HASH)
 app = FastAPI()
@@ -69,6 +70,30 @@ async def auth(action="check"):
 async def files(path=".") -> List:
     ret = []
 
+    base = Path(DB_PATH).joinpath("files")
+    path_obj = base.joinpath(path)
+
+    if not path_obj.exists():
+        return []
+
+    for file in path_obj.iterdir():
+        if file.is_dir():
+            ret.append({
+                "folder": True,
+                "name": file.name,
+                "path": str(file.relative_to(base)),
+                "id": ""
+            })
+
+    for line in path_obj.joinpath("files.txt").read_text().split("\n"):
+        file_id, file_name = line.split("~", 1)
+        ret.append({
+            "folder": False,
+            "name": file_name,
+            "path": str(path_obj.joinpath(file_name).relative_to(base)),
+            "id": file_id
+        })
+
     return ret
 
 
@@ -89,9 +114,31 @@ async def local_files(path="") -> List:
 async def tasks(request: Request, background_tasks: BackgroundTasks):
     body = await request.json()
 
-    # background_tasks.add_task(
+    if "source" not in body or "destination" not in body:
+        raise HTTPException(status_code=400, detail="Source and destination are required")
+
+    if "action" not in body:
+        raise HTTPException(status_code=400, detail="Action is required")
+
+    if body["action"] not in ["upload", "sync"]:
+        raise HTTPException(status_code=400, detail="Unknown action")
+
+    task = Task(f"{body['action'].title()} {body['source']}")
+    task.save()
+
+    functions = {
+        "upload": upload,
+        "sync": sync
+    }
+
+    background_tasks.add_task(
+        functions[body["action"]],
+        task,
+        Path(body["source"]),
+        Path(body["destination"])
+    )
 
-    return body
+    return task.signature()
 
 
 @app.get("/api/tasks/")

gramup/models.py

@@ -37,12 +37,11 @@ def to_dict(self):
             "message": self.message,
         }
 
-    def set_status(self, status):
-        self.status = status
-        self.save()
+    def set(self, *, status=None, message=None):
+        self.status = status or self.status
+        if message is not None:
+            self.message += message
 
-    def set_message(self, message):
-        self.message = message
         self.save()
 
     def load(self, task_id, db_path=DB_PATH):

gramup/tasks.py

@@ -1,3 +1,4 @@
+import asyncio
 from pathlib import Path
 
 from telethon import TelegramClient
@@ -7,15 +8,22 @@
 from gramup.models import Task
 
 
-async def pull_all_to_db(task: Task):
+async def get_client(task):
     client = TelegramClient("anon", API_ID, API_HASH)
 
     await client.connect()
 
     if not await client.is_user_authorized():
-        task.status = "failed"
-        task.message = "User not authorized"
-        task.save()
+        task.set(status="failed", message="Not authorized")
+        return
+
+    return client
+
+
+async def pull_all_to_db(task: Task):
+    client = await get_client(task)
+
+    if not client:
         return
 
     try:
@@ -25,24 +33,60 @@ async def pull_all_to_db(task: Task):
             caption: str = message.message
             path = Path(caption)
             name = path.name
-            path = path if path.parent != "." else path.joinpath("external")
             path = str(Path(DB_PATH).joinpath("files", path).with_name("files.txt"))
 
             if path not in to_write:
                 to_write[path] = []
 
-            to_write[path].append(" ".join((m_id, name)))
+            to_write[path].append(f"{m_id}~{name}")
 
         for path, messages in to_write.items():
             path = Path(path)
             path.parent.mkdir(parents=True, exist_ok=True)
-            path.write_text("\n".join(messages))
+            path.write_text("\n".join(messages), encoding='utf-8')
+
+    except Exception as e:
+        task.set(status="failed", message=f"Something went wrong: {e}")
+
+    task.set(status="completed")
+
 
+async def upload(task: Task, source: Path, destination: Path, branch=False):
+    client = await get_client(task)
+
+    if not client:
+        return
+
+    try:
+        sources = []
+
+        if source.is_dir():
+            tasks = []
+            for file in source.iterdir():
+                if file.is_file():
+                    sources.append(file)
+                else:
+                    tasks.append(upload(task, file, destination.joinpath(file.name), branch=True))
+
+            status = await asyncio.gather(*tasks, return_exceptions=True)
+            if any(status):
+                task.set(status="failed", message=f"Something went wrong: {status}")
+
+        else:
+            sources = [source]
+
+        tasks = []
+
+        for source in sources:
+            caption = str(destination.joinpath(source.name))
+            tasks.append(client.send_file("me", str(source), caption=caption, force_document=True))
+
+        await asyncio.gather(*tasks, return_exceptions=True)
     except Exception as e:
-        task.status = "failed"
-        task.message += f"Something went wrong: {e}"
-        task.save()
-        raise e
+        task.set(status="failed", message=f"Something went wrong: {e}")
+        return
 
-    task.status = "completed"
-    task.save()
+    if not branch:
+        task.set(status="completed")
+    else:
+        task.set(message=f"Uploaded {source.name}")