DRMDTool/main.go

package main

import (
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"embed"

	"github.com/BurntSushi/toml"
)

type Item struct {
	MPD         string
	Keys        string
	Filename    string
	Description string
	Subtitles   string
	Poster      string
	Metadata    string
}

type Items struct {
	Items []Item
}

type Metadata struct {
	Title  string
	Type   string
	Season string
}

type Config struct {
	BaseDir    string
	Format     string
	N_m3u8DLRE struct {
		Path string
	}
}

var config Config
var progressMutex sync.Mutex
var progress = make(map[string]*ProgressInfo)

const uploadDir = "uploads"

type ProgressInfo struct {
	Percentage  float64
	CurrentFile string
}

var templates *template.Template

var templateFS embed.FS

func init() {
	if err := os.MkdirAll(uploadDir, 0755); err != nil {
		fmt.Printf("Error creating upload directory: %v\n", err)
	}

	templates = template.Must(template.ParseFS(templateFS, "templates/*"))
}

func main() {
	configFile, err := os.Open("config.toml")
	if err != nil {
		fmt.Println("Error opening config file:", err)
		return
	}
	defer configFile.Close()

	byteValue, _ := io.ReadAll(configFile)

	if _, err := toml.Decode(string(byteValue), &config); err != nil {
		fmt.Println("Error decoding config file:", err)
		return
	}

	if config.N_m3u8DLRE.Path == "" {
		fmt.Println("Error: N_m3u8DL-RE path is not specified in the config file")
		return
	}

	inputFile := flag.String("f", "", "Path to the input JSON file")
	flag.Parse()

	if *inputFile == "" {
		startWebServer()
	} else {
		processInputFile(*inputFile)
	}
}

func startWebServer() {
	http.HandleFunc("/", handleRoot)
	http.HandleFunc("/upload", handleUpload)
	http.HandleFunc("/progress", handleProgress)

	fmt.Println("Starting web server on http://0.0.0.0:8080")
	http.ListenAndServe(":8080", nil)
}

func handleRoot(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}

	progressMutex.Lock()
	jobs := make(map[string]*ProgressInfo)
	for k, v := range progress {
		jobs[k] = v
	}
	progressMutex.Unlock()

	err := templates.ExecuteTemplate(w, "index", struct{ Jobs map[string]*ProgressInfo }{jobs})
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func handleUpload(w http.ResponseWriter, r *http.Request) {
	file, header, err := r.FormFile("file")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	defer file.Close()

	filename := fmt.Sprintf("%d_%s", time.Now().UnixNano(), header.Filename)
	filepath := filepath.Join(uploadDir, filename)

	newFile, err := os.Create(filepath)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer newFile.Close()

	_, err = io.Copy(newFile, file)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	go func() {
		err := processInputFile(filepath)
		if err != nil {
			fmt.Printf("Error processing file: %v\n", err)
		}

		os.Remove(filepath)
	}()

	http.Redirect(w, r, "/progress?filename="+filename, http.StatusSeeOther)
}

func handleProgress(w http.ResponseWriter, r *http.Request) {
	filename := r.URL.Query().Get("filename")
	fmt.Printf("Handling progress request for filename: %s\n", filename)

	if r.Header.Get("Accept") == "application/json" {
		progressInfo := getProgress(filename)
		fmt.Printf("Progress info for %s: %+v\n", filename, progressInfo)

		if progressInfo == nil {
			w.WriteHeader(http.StatusNotFound)
			json.NewEncoder(w).Encode(map[string]string{"error": "No progress information found"})
			return
		}

		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(progressInfo)
		if err != nil {
			fmt.Printf("Error encoding progress info: %v\n", err)
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}
		return
	}

	err := templates.ExecuteTemplate(w, "progress", struct{ Filename string }{filename})
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func processInputFile(inputFile string) error {
	jsonFile, err := os.Open(inputFile)
	if err != nil {
		return fmt.Errorf("error opening file %s: %v", inputFile, err)
	}
	defer jsonFile.Close()

	byteValue, err := io.ReadAll(jsonFile)
	if err != nil {
		return fmt.Errorf("error reading file %s: %v", inputFile, err)
	}

	var items Items
	err = json.Unmarshal(byteValue, &items)
	if err != nil {
		return fmt.Errorf("error unmarshaling JSON: %v", err)
	}

	for i, item := range items.Items {
		updateProgress(filepath.Base(inputFile), float64(i)/float64(len(items.Items))*100, item.Filename)
		err := downloadFile(item)
		if err != nil {
			fmt.Printf("Error downloading file: %v\n", err)
		}
	}
	updateProgress(filepath.Base(inputFile), 100, "")

	return nil
}

func updateProgress(filename string, value float64, currentFile string) {
	progressMutex.Lock()
	defer progressMutex.Unlock()
	progress[filename] = &ProgressInfo{
		Percentage:  value,
		CurrentFile: currentFile,
	}
	fmt.Printf("Progress updated for %s: %.2f%%, Current file: %s\n", filename, value, currentFile)
}

func getProgress(filename string) *ProgressInfo {
	progressMutex.Lock()
	defer progressMutex.Unlock()
	return progress[filename]
}

func getKeys(keys string) []string {
	return strings.Split(keys, ",")
}

func parseMetadata(metadata string) Metadata {
	parts := strings.Split(metadata, ";")
	if len(parts) != 3 {
		return Metadata{}
	}
	return Metadata{
		Title:  strings.TrimSpace(parts[0]),
		Type:   strings.TrimSpace(parts[1]),
		Season: "S" + strings.TrimSpace(parts[2]),
	}
}

func getDownloadCommand(item Item, mpdPath string) string {
	metadata := parseMetadata(item.Metadata)
	keys := getKeys(item.Keys)

	command := fmt.Sprintf("%s %s", config.N_m3u8DLRE.Path, mpdPath)

	for _, key := range keys {
		if key != "" {
			command += fmt.Sprintf(" --key %s", key)
		}
	}

	command += " --auto-select"

	filename := fmt.Sprintf("\"%s\"", item.Filename)
	command += fmt.Sprintf(" --save-name %s", filename)

	command += fmt.Sprintf(" --mux-after-done format=%s", config.Format)

	saveDir := config.BaseDir
	if metadata.Type == "serie" {
		saveDir = filepath.Join(saveDir, "Series", metadata.Title, metadata.Season)
	} else {
		saveDir = filepath.Join(saveDir, "Movies", metadata.Title)
	}
	command += fmt.Sprintf(" --save-dir \"%s\"", saveDir)

	fmt.Println(command)

	return command
}

func downloadFile(item Item) error {
	fmt.Println("Downloading:", item.Filename)

	mpdPath := item.MPD
	if !isValidURL(item.MPD) {

		decodedMPD, err := base64.StdEncoding.DecodeString(item.MPD)
		if err != nil {
			return fmt.Errorf("error decoding base64 MPD: %v", err)
		}

		tempFile, err := os.CreateTemp("", "temp_mpd_*.mpd")
		if err != nil {
			return fmt.Errorf("error creating temporary MPD file: %v", err)
		}
		defer os.Remove(tempFile.Name())

		if _, err := tempFile.Write(decodedMPD); err != nil {
			return fmt.Errorf("error writing to temporary MPD file: %v", err)
		}
		if err := tempFile.Close(); err != nil {
			return fmt.Errorf("error closing temporary MPD file: %v", err)
		}

		mpdPath = tempFile.Name()
	}

	command := getDownloadCommand(item, mpdPath)

	cmd := exec.Command("bash", "-c", command)

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("error executing download command: %v", err)
	}

	fmt.Println("Download completed successfully")
	return nil
}

func isValidURL(toTest string) bool {
	_, err := url.ParseRequestURI(toTest)
	return err == nil
}