Add encryption, download files

.gitignore

@@ -1,7 +1,9 @@
-partstostitch/aduhm.mp4
-partstostitch/aduhm.part1
-partstostitch/aduhm.part2
-partstostitch/aduhm.part3
-partstostitch/aduhm.part4
-partstostitch/aduhm.part5
-partstostitch/aduhm.part6
+partstostitch
+.env
+__pycache__/
+.idea/
+venv/
+.vscode/
+db.json
+.DS_Store
+**/.DS_Store

README.md

@@ -13,7 +13,14 @@ To deploy this project you'll need to install the dependencies included in the r
   pip install -r requirements.txt
 ```
 
-Replace the client token and the channel ID in app.py with your client token and your channel ID. If you don't have a Discord bot made and invited to a server already, you will need to do so.
+If you don't have a Discord bot made and invited to a server already, you will need to do so.
+
+Create a file called `.env` and add the following to it:
+
+```env
+DISCORD_TOKEN=your_discord_token
+DISCORD_CHANNEL=your_discord_channel_id
+```
 
 Following that, run
 

app.py

@@ -1,49 +1,183 @@
-from flask import Flask, request, jsonify, render_template
+from flask import Flask, request, jsonify, render_template, send_from_directory
 import discord
 import asyncio
 import threading
+import io
+from dotenv import load_dotenv
+import os
+import json
+
+load_dotenv()
 
 app = Flask(__name__)
 
 # Initialize Discord client outside the request handling
 intents = discord.Intents.default()
 client = discord.Client(intents=intents)
 
-@app.route('/')
+
+@app.route("/")
 def index():
-    return render_template('index.html')
+    return render_template("index.html")
+
+
+@app.route("/upload")
+def upload():
+    return render_template("upload.html")
+
+
+@app.route("/download")
+def download():
+    # Read the existing data from db.json
+    try:
+        with open("db.json", "r") as db_file:
+            db_data = json.load(db_file)
+            files = db_data.get("files", {})
+    except (FileNotFoundError, json.JSONDecodeError):
+        files = {}
+
+    # Convert the files dictionary to a list including the number of parts
+    files_list = [name for name, info in files.items()]
+
+    return render_template("download.html", files=files_list)
+
 
 def start_discord_bot():
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
 
     @client.event
     async def on_ready():
-        print('Discord client is ready.')
+        print("Discord client is ready.")
+
+    client.run(os.environ["DISCORD_TOKEN"])
 
-    client.run('') #Replace with client token
 
 # Run Discord client in a background thread
 threading.Thread(target=start_discord_bot, daemon=True).start()
 
-@app.route('/upload', methods=['POST'])
+
+@app.route("/uploadFile", methods=["POST"])
 def upload_file():
-    file = request.files['file']
-
+    file = request.files["file"]
+    filename = file.filename
+
     # Ensure the Discord client is ready before attempting to send
     if not client.is_closed():
-        future = asyncio.run_coroutine_threadsafe(send_file_to_discord(file), client.loop)
-        result = future.result()  # This waits for the coroutine to finish and returns its result
+        future = asyncio.run_coroutine_threadsafe(
+            send_file_to_discord(file), client.loop
+        )
+        result = (
+            future.result()
+        )  # This waits for the coroutine to finish and returns its result
         print(result)
 
-
-    return jsonify({'message': 'File uploaded successfully!'})
+    # Read the existing data from db.json
+    try:
+        with open("db.json", "r") as db_file:
+            db_data = json.load(db_file)
+    except (FileNotFoundError, json.JSONDecodeError):
+        db_data = {"files": {}}
+
+    # Increment the parts count for the uploaded file
+    if filename in db_data["files"]:
+        db_data["files"][filename.split(".part")[0]]["parts"] += 1
+    else:
+        db_data["files"][filename.split(".part")[0]] = {"parts": 1}
+
+    # Write the updated data back to db.json
+    with open("db.json", "w") as db_file:
+        json.dump(db_data, db_file, indent=4)
+
+    return jsonify({"message": "File uploaded successfully!"})
+
+
+should_remove_files = False
+
+
+@app.after_request
+def remove_files(response):
+    global should_remove_files
+    if should_remove_files:
+        for file in os.listdir("partstostitch"):
+            os.remove(os.path.join("partstostitch", file))
+        should_remove_files = False
+    return response
+
+
+@app.route("/downloadFile/<filename>")
+def download_file(filename):
+    global should_remove_files
+    # Ensure the directory exists
+    os.makedirs("partstostitch", exist_ok=True)
+
+    # Ensure the Discord client is ready before attempting to fetch messages
+    if not client.is_closed():
+        # Run the coroutine in the same thread as the Discord client's event loop
+        future = asyncio.run_coroutine_threadsafe(
+            fetch_and_save_files(filename), client.loop
+        )
+        result = (
+            future.result()
+        )  # This waits for the coroutine to finish and returns its result
+        if result > 1:  # Check if multiple files were downloaded
+            # Merge the files
+            with open(f"partstostitch/{filename}", "wb") as merged_file:
+                for part_num in range(1, result + 1):
+                    part_filename = f"{filename}.part{part_num}"
+                    part_path = os.path.join("partstostitch", part_filename)
+                    with open(part_path, "rb") as part_file:
+                        merged_file.write(part_file.read())
+                    os.remove(part_path)  # Remove the part file after merging
+            should_remove_files = True
+            return send_from_directory(
+                "partstostitch",
+                filename,
+                as_attachment=True,
+                mimetype="application/octet-stream",
+            )
+        elif result == 1:
+            # Only one file, no need to merge
+            should_remove_files = True
+            return send_from_directory(
+                "partstostitch",
+                filename,
+                as_attachment=True,
+                mimetype="application/octet-stream",
+            )
+        else:
+            # No files were downloaded
+            return jsonify({"error": "No files were found to download"}), 404
+    else:
+        return jsonify({"error": "Discord client is not ready"}), 503
+
+
+async def fetch_and_save_files(filename):
+    channel = client.get_channel(int(os.environ["DISCORD_CHANNEL"]))
+    messages = [message async for message in channel.history()]
+    count = 0
+
+    for message in messages:
+        for attachment in message.attachments:
+            if attachment.filename.startswith(filename):
+                file_bytes = await attachment.read()
+                file_path = os.path.join("partstostitch", attachment.filename)
+                with open(file_path, "wb") as file:
+                    file.write(file_bytes)
+                count += 1
+
+    return count
+
 
 async def send_file_to_discord(file):
     print("Sending file to Discord:", file.filename)
-    channel = client.get_channel()  # Replace with channel ID
-    message = await channel.send(file=discord.File(file.stream, filename=file.filename))
+    channel = client.get_channel(int(os.environ["DISCORD_CHANNEL"]))
+    file_bytes = file.read()
+    message = await channel.send(
+        file=discord.File(io.BytesIO(file_bytes), filename=file.filename)
+    )
     print(message)
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     app.run(debug=True)

requirements.txt

@@ -22,6 +22,7 @@ MarkupSafe==2.1.5
 mdurl==0.1.2
 multidict==6.0.5
 Pygments==2.17.2
+python-dotenv==1.0.1
 python-nmap==0.7.1
 requests==2.31.0
 rich==13.7.0