import argparse
import json
import sys
import shutil
import pathlib
import random
import re
import requests
from os import path

ap = argparse.ArgumentParser(description = "Downloads attachments and image embeds from a DHT archive into one folder, and updates URLs in the archive to point to the downloaded files. The original archive will be backed up with a '.bak' extension; if a backup file already exists, it will NOT be overwritten.")

ap.add_argument("input_archive",
                metavar = "input-archive",
                help = "path to archive file")

ap.add_argument("-d", "--download-folder-name",
                default = "",
                required = False,
                metavar = "name",
                dest = "download_folder_name",
                help = "name of folder with downloaded files (if omitted, will be same as input archive name with '.downloads' appended to the end)")

if len(sys.argv) == 1:
    ap.print_help()
    exit(1)

args = ap.parse_args()

input_archive = args.input_archive
download_folder_name = args.download_folder_name if len(args.download_folder_name) else path.basename(input_archive) + ".downloads"
download_folder_path = path.join(pathlib.Path(input_archive).resolve().with_name(download_folder_name))

# Setup

try:
    pathlib.Path(download_folder_path).mkdir(parents = True, exist_ok = True)
except OSError as e:
    print("Could not create download folder: " + str(e))
    exit(1)

try:
    with open(input_archive, "r", encoding = "UTF-8") as f:
        archive = json.load(f)
except FileNotFoundError:
    print("Input archive file not found: " + input_archive)
    exit(1)
except json.JSONDecodeError as e:
    print("Input archive file has invalid format: " + str(e))
    exit(1)

backup_archive = input_archive + ".bak"

if not path.isfile(backup_archive):
    shutil.copy(input_archive, backup_archive)

# Collect

download_url_to_ele = {}
download_url_to_file = {}
download_file_to_url = {}


def add_file_to_download(ele):
    url = ele["url"]
    
    if url.startswith("file:"):
        return
    
    url_split = url.split("://", maxsplit = 2)
    
    if len(url_split) != 2:
        print("Invalid attachment URL: " + url)
        return
    
    download_file_name = re.sub(r"[^\w\-_.]", "_", url_split[1])
    
    while download_file_name in download_file_to_url:
        download_file_name += "_" + str(random.randint(0, 9))
    
    download_url_to_ele[url] = ele
    download_url_to_file[url] = download_file_name
    download_file_to_url[download_file_name] = url


archive_data = archive["data"]

for channel_id, channel_data in archive_data.items():
    for message_id, message_data in channel_data.items():
        if "a" in message_data:
            for attachment in message_data["a"]:
                add_file_to_download(attachment)
        
        if "e" in message_data:
            for embed in message_data["e"]:
                if embed["type"] == "image":
                    add_file_to_download(embed)

# Download

counter = 0
total = len(download_url_to_file)
digits = len(str(total))

print("Identified {} attachment(s) and image embed(s) to download.".format(total))
print("")

failed = list()

for url, file in download_url_to_file.items():
    counter += 1
    print("[{}/{}] {}".format(str(counter).rjust(digits, " "), total, url))
    
    full_path = download_folder_path + "/" + file
    
    if path.isfile(full_path):
        print("Already downloaded, skipping...")
    else:
        try:
            req = requests.get(url, timeout = 2)
            req.raise_for_status()
            
            with open(full_path, "wb") as f:
                f.write(req.content)
        
        except Exception as e:
            failed.append((url, file))
            print("Download failed... {}".format(e))
            continue
    
    download_url_to_ele[url]["url"] = "file:./" + download_folder_name + "/" + file

# Update

with open(input_archive, "w", encoding = "UTF-8") as f:
    json.dump(archive, f, separators = (",", ":"))

print("")

if len(failed) > 0:
    print("Archive was updated, but {} out of {} request(s) failed. You may re-run the script to try re-downloading failed requests again.".format(len(failed), total))
else:
    print("Archive was updated.")

print("To view the archive with downloaded files, you must place the viewer in the same folder as the archive file and download folder.")

Add Python script to download attachments and image embeds & update the archive