Merge branch 'main' into speedLimiter

Reviewed-on: #9

README.md

@@ -7,17 +7,51 @@ drmdtool is a utility for processing .drmd files using N_m3u8DL-RE.
 Create a `config.toml` file in the same directory as the drmdtool executable:
 
 ```toml
+[General]
 BaseDir = "/path/to/save/downloads"
 Format = "mkv"
 TempBaseDir = "/tmp/nre"
 EnableConsole = true
-WatchedFolder = "/path/to/watched/folder"
 
-[N_m3u8DL-RE]
+[WatchFolder]
+Path = "/path/to/watched/folder"
+PollingInterval = 10
+UsePolling = true
+UseInotify = false
+
+[N_m3u8DLRE]
 Path = "/path/to/N_m3u8DL-RE"
 ```
 
-Adjust the paths and format as needed. (mkv, mp4)
+### Configuration Options
+
+- **General**
+  - `BaseDir`: Directory where downloaded files will be saved.
+  - `Format`: Output format for the downloaded files (e.g., `mkv`, `mp4`).
+  - `TempBaseDir`: Temporary directory for intermediate files.
+  - `EnableConsole`: Boolean to enable or disable console output.
+
+- **WatchFolder**
+  - `Path`: Directory to watch for new `.drmd` files.
+  - `PollingInterval`: Interval in seconds for polling the watch folder.
+  - `UsePolling`: Boolean to enable or disable folder polling.
+  - `UseInotify`: Boolean to enable or disable inotify for file watching.
+
+- **N_m3u8DLRE**
+  - `Path`: Path to the N_m3u8DL-RE executable.
+
+### Environment Variable Overrides
+
+You can override the configuration options using environment variables. The following environment variables are supported:
+
+- `BASE_DIR`: Overrides `General.BaseDir`
+- `FORMAT`: Overrides `General.Format`
+- `TEMP_BASE_DIR`: Overrides `General.TempBaseDir`
+- `ENABLE_CONSOLE`: Overrides `General.EnableConsole` (set to `true` or `false`)
+- `WATCHED_FOLDER`: Overrides `WatchFolder.Path`
+- `USE_POLLING`: Overrides `WatchFolder.UsePolling` (set to `true` or `false`)
+- `USE_INOTIFY`: Overrides `WatchFolder.UseInotify` (set to `true` or `false`)
+- `POLLING_INTERVAL`: Overrides `WatchFolder.PollingInterval`
 
 ## Web UI Usage
 
@@ -30,7 +64,6 @@ Adjust the paths and format as needed. (mkv, mp4)
 
 3. Use the interface to upload .drmd files and monitor download progress
 
-
 ## CLI Usage
 
 To process a file directly from the command line:
@@ -41,7 +74,6 @@ To process a file directly from the command line:
 
 This will download the file and save it in the base directory specified in the config.
 
-
 # Previews
 
 ## Index Page

config.toml

@@ -1,8 +1,14 @@
-BaseDir = "/mnt/media"
+[General]
+BaseDir = "/mnt/Media"
 Format = "mkv"
 TempBaseDir = "/tmp/nre"
 EnableConsole = true
-WatchedFolder = "/mnt/watched"
+
+[WatchFolder]
+Path = "/mnt/Watched"
+PollingInterval = 10
+UsePolling = false 
+UseInotify = false 
 
 [N_m3u8DLRE]
 Path = "nre"

src/config.go

@@ -4,19 +4,28 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strconv"
+	"strings"
 
 	"github.com/BurntSushi/toml"
 )
 
 type Config struct {
-	BaseDir     string
-	Format      string
-	TempBaseDir string
-	N_m3u8DLRE  struct {
+	General struct {
+		BaseDir       string
+		Format        string
+		TempBaseDir   string
+		EnableConsole bool
+	}
+	WatchFolder struct {
+		Path            string
+		UsePolling      bool
+		UseInotify      bool
+		PollingInterval int
+	}
+	N_m3u8DLRE struct {
 		Path string
 	}
-	EnableConsole bool
-	WatchedFolder string
 }
 
 var config Config
@@ -24,25 +33,112 @@ var config Config
 func loadConfig() {
 	configFile, err := os.Open("config.toml")
 	if err != nil {
-		fmt.Println("Error opening config file:", err)
-		return
+		logger.LogError("Config", fmt.Sprintf("Error opening config file: %v", err))
+		os.Exit(1)
 	}
 	defer configFile.Close()
 
 	byteValue, _ := io.ReadAll(configFile)
 
 	if _, err := toml.Decode(string(byteValue), &config); err != nil {
-		fmt.Println("Error decoding config file:", err)
-		return
+		logger.LogError("Config", fmt.Sprintf("Error decoding config file: %v", err))
+		os.Exit(1)
 	}
 
-	if config.N_m3u8DLRE.Path == "" {
-		fmt.Println("Error: N_m3u8DL-RE path is not specified in the config file")
-		return
+	overrideConfigWithEnv()
+
+	if err := validatePaths(); err != nil {
+		logger.LogError("Config", fmt.Sprintf("Configuration error: %v", err))
+		os.Exit(1)
 	}
 
-	if config.WatchedFolder == "" {
-		fmt.Println("Error: Watched folder is not specified in the config file")
-		return
+	if config.WatchFolder.PollingInterval <= 0 {
+		config.WatchFolder.PollingInterval = 10
+	}
+
+	logConfig()
+}
+
+func overrideConfigWithEnv() {
+	if envBaseDir := os.Getenv("BASE_DIR"); envBaseDir != "" {
+		config.General.BaseDir = envBaseDir
+	}
+	if envFormat := os.Getenv("FORMAT"); envFormat != "" {
+		config.General.Format = envFormat
+	}
+	if envTempBaseDir := os.Getenv("TEMP_BASE_DIR"); envTempBaseDir != "" {
+		config.General.TempBaseDir = envTempBaseDir
+	}
+	if envEnableConsole := os.Getenv("ENABLE_CONSOLE"); envEnableConsole != "" {
+		config.General.EnableConsole = strings.ToLower(envEnableConsole) == "true"
+	}
+	if envWatchedFolder := os.Getenv("WATCHED_FOLDER"); envWatchedFolder != "" {
+		config.WatchFolder.Path = envWatchedFolder
+	}
+	if envUsePolling := os.Getenv("USE_POLLING"); envUsePolling != "" {
+		config.WatchFolder.UsePolling = strings.ToLower(envUsePolling) == "true"
+	}
+	if envUseInotify := os.Getenv("USE_INOTIFY"); envUseInotify != "" {
+		config.WatchFolder.UseInotify = strings.ToLower(envUseInotify) == "true"
+	}
+	if envPollingInterval := os.Getenv("POLLING_INTERVAL"); envPollingInterval != "" {
+		if interval, err := strconv.Atoi(envPollingInterval); err == nil {
+			config.WatchFolder.PollingInterval = interval
+		}
 	}
 }
+
+func validatePaths() error {
+	paths := []struct {
+		name string
+		path string
+	}{
+		{"BaseDir", config.General.BaseDir},
+	}
+
+	for _, p := range paths {
+		if p.path == "" {
+			return fmt.Errorf("%s is not specified", p.name)
+		}
+		if _, err := os.Stat(p.path); os.IsNotExist(err) {
+			return fmt.Errorf("%s does not exist: %s", p.name, p.path)
+		} else if err != nil {
+			return fmt.Errorf("error accessing %s: %v", p.name, err)
+		}
+	}
+
+	if config.WatchFolder.UsePolling || config.WatchFolder.UseInotify {
+		if config.WatchFolder.Path == "" {
+			return fmt.Errorf("WatchedFolder is not specified")
+		}
+		if _, err := os.Stat(config.WatchFolder.Path); os.IsNotExist(err) {
+			return fmt.Errorf("WatchedFolder does not exist: %s", config.WatchFolder.Path)
+		} else if err != nil {
+			return fmt.Errorf("error accessing WatchedFolder: %v", err)
+		}
+	}
+
+	return nil
+}
+
+func logConfig() {
+	configInfo := fmt.Sprintf(`
+Configuration Loaded:
+  General:
+    BaseDir: %s
+    Format: %s
+    TempBaseDir: %s
+    EnableConsole: %t
+  WatchFolder:
+    Path: %s
+    UsePolling: %t
+    UseInotify: %t
+    PollingInterval: %d
+  N_m3u8DLRE:
+    Path: %s
+`, config.General.BaseDir, config.General.Format, config.General.TempBaseDir, config.General.EnableConsole,
+		config.WatchFolder.Path, config.WatchFolder.UsePolling, config.WatchFolder.UseInotify, config.WatchFolder.PollingInterval,
+		config.N_m3u8DLRE.Path)
+
+	logger.LogInfo("Config", configInfo)
+}

src/downloaders.go

@@ -23,7 +23,7 @@ func removeBOM(input []byte) []byte {
 func downloadFile(drmdFilename string, item Item, jobInfo *JobInfo) error {
 	logger.LogInfo("Download File", fmt.Sprintf("Starting download for: %s", item.Filename))
 
-	tempDir := filepath.Join(config.TempBaseDir, sanitizeFilename(item.Filename))
+	tempDir := filepath.Join(config.General.TempBaseDir, sanitizeFilename(item.Filename))
 	err := os.MkdirAll(tempDir, 0755)
 	if err != nil {
 		logger.LogError("Download File", fmt.Sprintf("Error creating temporary directory: %v", err))
@@ -132,7 +132,7 @@ func downloadFile(drmdFilename string, item Item, jobInfo *JobInfo) error {
 		for {
 			if outputBuffer.Len() > 0 {
 				message := outputBuffer.Bytes()
-				if config.EnableConsole {
+				if config.General.EnableConsole {
 					broadcast(drmdFilename, message)
 				}
 				outputBuffer.Reset()
@@ -183,9 +183,9 @@ func getDownloadCommand(item Item, mpdPath string, tempDir string) string {
 	filename := fmt.Sprintf("\"%s\"", sanitizedFilename)
 	command += fmt.Sprintf(" --save-name %s", filename)
 
-	command += fmt.Sprintf(" --mux-after-done format=%s", config.Format)
+	command += fmt.Sprintf(" --mux-after-done format=%s", config.General.Format)
 
-	saveDir := config.BaseDir
+	saveDir := config.General.BaseDir
 	if metadata.Type == "serie" {
 		saveDir = filepath.Join(saveDir, "Series", metadata.Title, metadata.Season)
 	} else {
@@ -195,6 +195,10 @@ func getDownloadCommand(item Item, mpdPath string, tempDir string) string {
 
 	command += fmt.Sprintf(" --tmp-dir \"%s\"", tempDir)
 
+	if globalSpeedLimit != "" {
+		command += fmt.Sprintf(" -R %s", globalSpeedLimit)
+	}
+
 	fmt.Println(command)
 
 	return command

src/handlers.go

@@ -48,7 +48,11 @@ func handleRoot(w http.ResponseWriter, r *http.Request) {
 			CurrentFile string
 			Paused      bool
 		}
-	}{jobsInfo})
+		GlobalSpeedLimit string
+	}{
+		Jobs:             jobsInfo,
+		GlobalSpeedLimit: globalSpeedLimit,
+	})
 	if err != nil {
 		logger.LogError("Handle Root", fmt.Sprintf("Error executing template: %v", err))
 		http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -386,8 +390,8 @@ var clients = make(map[string]map[*websocket.Conn]bool)
 var mu sync.Mutex
 
 func handleWebSocket(w http.ResponseWriter, r *http.Request) {
-	fmt.Println(config.EnableConsole)
-	if !config.EnableConsole {
+	fmt.Println(config.General.EnableConsole)
+	if !config.General.EnableConsole {
 		http.Error(w, "Console output is disabled", http.StatusForbidden)
 		return
 	}
@@ -428,7 +432,7 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request) {
 }
 
 func broadcast(filename string, message []byte) {
-	if !config.EnableConsole {
+	if !config.General.EnableConsole {
 		return
 	}
 
@@ -443,3 +447,32 @@ func broadcast(filename string, message []byte) {
 		}
 	}
 }
+
+func handleSetSpeedLimit(w http.ResponseWriter, r *http.Request) {
+	logger.LogInfo("Set Speed Limit", "Received request to set speed limit")
+
+	if r.Method != http.MethodPost {
+		logger.LogError("Set Speed Limit", "Invalid method")
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	var requestData struct {
+		SpeedLimit string `json:"speedLimit"`
+	}
+
+	if err := json.NewDecoder(r.Body).Decode(&requestData); err != nil {
+		logger.LogError("Set Speed Limit", "Invalid request body")
+		http.Error(w, "Invalid request", http.StatusBadRequest)
+		return
+	}
+
+	if requestData.SpeedLimit == "unlimited" {
+		globalSpeedLimit = ""
+	} else {
+		globalSpeedLimit = requestData.SpeedLimit
+	}
+
+	logger.LogInfo("Set Speed Limit", fmt.Sprintf("Global speed limit set to: %s", globalSpeedLimit))
+	w.WriteHeader(http.StatusOK)
+}

src/main.go

@@ -44,6 +44,8 @@ var templates *template.Template
 //go:embed templates
 var templateFS embed.FS
 
+var globalSpeedLimit string
+
 func init() {
 	if err := os.MkdirAll(uploadDir, 0755); err != nil {
 		fmt.Printf("Error creating upload directory: %v\n", err)
@@ -70,6 +72,8 @@ func main() {
 		}
 		processItems(*inputFile, items)
 	}
+
+	http.HandleFunc("/set-speed-limit", handleSetSpeedLimit)
 }
 
 func startWebServer() {
@@ -83,8 +87,9 @@ func startWebServer() {
 	http.HandleFunc("/resume", handleResume)
 	http.HandleFunc("/clear-completed", handleClearCompleted)
 	http.HandleFunc("/ws", handleWebSocket)
+	http.HandleFunc("/set-speed-limit", handleSetSpeedLimit)
 
-	fmt.Println("Starting web server on http://0.0.0.0:8080")
+	logger.LogInfo("Main", "Starting web server on http://0.0.0.0:8080")
 	http.ListenAndServe(":8080", nil)
 }
 

src/templates/index

@@ -28,7 +28,6 @@
         input[type="file"], input[type="submit"] {
             background-color: #2d2d2d;
             color: #d4d4d4;
-            border: 1px solid #444;
             padding: 8px 12px;
             border-radius: 4px;
             margin-bottom: 10px;
@@ -77,6 +76,11 @@
         .paused {
             color: #ffa500;
         }
+        .speed-limit {
+            font-size: 1em;
+            color: #a0a0a0;
+            margin-top: 10px;
+        }
         @media (max-width: 600px) {
             body {
                 padding: 10px;
@@ -107,6 +111,94 @@
         #clear-completed:hover {
             background-color: #d32f2f;
         }
+        /* New CSS for speed limit form */
+        .settings-section {
+            margin-top: 30px;
+        }
+        .speed-limit-form {
+            display: flex;
+            align-items: center;
+            justify-content: space-between;
+            gap: 10px;
+            margin-bottom: 20px;
+        }
+        .speed-limit-form .form-group {
+            display: flex;
+            align-items: center;
+            gap: 10px;
+        }
+        .speed-limit-form input[type="number"],
+        .speed-limit-form select,
+        .speed-limit-form button {
+            background-color: #2d2d2d;
+            color: #d4d4d4;
+            border: 1px solid #444;
+            padding: 8px 12px;
+            border-radius: 4px;
+        }
+        .speed-limit-form button {
+            cursor: pointer;
+            background-color: #4CAF50;
+            color: white;
+        }
+        .speed-limit-form button:hover {
+            background-color: #45a049;
+        }
+        .speed-limit-container {
+            display: flex;
+            align-items: center;
+            margin-bottom: 20px;
+            background-color: #2d2d2d;
+            padding: 8px 12px;
+            border-radius: 4px;
+        }
+        .speed-limit-container .form-group {
+            display: flex;
+            align-items: center;
+            gap: 10px;
+            width: 100%;
+        }
+        .speed-limit-container input[type="number"] {
+            background-color: #2d2d2d;
+            color: #d4d4d4;
+            border: 1px solid #444;
+            padding: 8px 12px;
+            border-radius: 4px;
+            height: 40px;
+            box-sizing: border-box;
+            flex-grow: 1;
+        }
+        .speed-limit-container select,
+        .speed-limit-container button {
+            background-color: #2d2d2d;
+            color: #d4d4d4;
+            border: 1px solid #444;
+            padding: 8px 12px;
+            border-radius: 4px;
+            height: 40px;
+            box-sizing: border-box;
+        }
+        .speed-limit-container button {
+            cursor: pointer;
+            background-color: #4CAF50;
+            color: white;
+        }
+        .speed-limit-container button:hover {
+            background-color: #45a049;
+        }
+        .speed-limit-container .speed-limit {
+            color: #d4d4d4;
+            margin-left: auto;
+            display: flex;
+            align-items: center;
+        }
+        .speed-limit-container .speed-limit span {
+            margin-left: 5px;
+        }
+        .current-speed-limit {
+            color: #d4d4d4;
+            margin-top: 10px;
+        }
     </style>
 </head>
 <body>
@@ -135,6 +227,23 @@
         {{end}}
     </ul>
     <button id="clear-completed" onclick="clearCompleted()">Clear Completed Jobs</button>
+    
+    <div class="settings-section">
+        <h2>Settings</h2>
+        <div class="speed-limit-container">
+            <div class="form-group">
+                <label for="speedLimitValue">Speed Limit:</label>
+                <input type="number" id="speedLimitValue" name="speedLimitValue" min="0" step="0.01" required>
+                <select id="speedLimitUnit" name="speedLimitUnit">
+                    <option value="GBps">GBps</option>
+                    <option value="MBps" selected>MBps</option>
+                    <option value="KBps">KBps</option>
+                </select>
+                <button type="button" onclick="updateSpeedLimit(event)">Set Limit</button>
+            </div>
+        </div>
+    </div>
+
     <script>
         function clearCompleted() {
             fetch('/clear-completed', { method: 'POST' })
@@ -147,6 +256,54 @@
                     }
                 });
         }
+
+        function updateSpeedLimit(event) {
+            event.preventDefault();
+
+            const speedLimitValue = document.getElementById('speedLimitValue').value;
+            const speedLimitUnit = document.getElementById('speedLimitUnit').value;
+            const speedLimit = speedLimitValue === "0" ? "unlimited" : speedLimitValue + speedLimitUnit;
+
+            if (!validateSpeedLimit(speedLimitValue)) {
+                alert('Please enter a valid speed limit.');
+                return;
+            }
+
+            fetch('/set-speed-limit', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({ speedLimit }),
+            }).then(response => {
+                if (response.ok) {
+                    alert('Speed limit updated successfully');
+                    document.getElementById('currentSpeedLimit').textContent = speedLimit;
+                } else {
+                    alert('Failed to update speed limit');
+                }
+            });
+        }
+
+        function validateSpeedLimit(value) {
+            const number = parseFloat(value);
+            return !isNaN(number) && number >= 0;
+        }
+
+        document.addEventListener('DOMContentLoaded', function() {
+            const currentSpeedLimit = "{{if .GlobalSpeedLimit}}{{.GlobalSpeedLimit}}{{else}}0{{end}}";
+            const speedLimitValueInput = document.getElementById('speedLimitValue');
+            const speedLimitUnitSelect = document.getElementById('speedLimitUnit');
+
+            const match = currentSpeedLimit.match(/(\d+(\.\d+)?)([A-Za-z]+)/);
+            if (match) {
+                speedLimitValueInput.value = match[1];
+                speedLimitUnitSelect.value = match[3];
+            } else {
+                speedLimitValueInput.value = "0"; 
+                speedLimitUnitSelect.value = "MBps";
+            }
+        });
     </script>
 </body>
 </html>

src/watcher.go

@@ -13,6 +13,16 @@ import (
 )
 
 func watchFolder() {
+	if config.WatchFolder.UsePolling {
+		go pollFolder()
+	}
+
+	if config.WatchFolder.UseInotify {
+		go inotifyWatch()
+	}
+}
+
+func inotifyWatch() {
 	watcher, err := fsnotify.NewWatcher()
 	if err != nil {
 		log.Fatal(err)
@@ -43,13 +53,31 @@ func watchFolder() {
 		}
 	}()
 
-	err = watcher.Add(config.WatchedFolder)
+	err = watcher.Add(config.WatchFolder.Path)
 	if err != nil {
 		log.Fatal(err)
 	}
 	<-done
 }
 
+func pollFolder() {
+	ticker := time.NewTicker(time.Duration(config.WatchFolder.PollingInterval) * time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		files, err := filepath.Glob(filepath.Join(config.WatchFolder.Path, "*.drmd"))
+		if err != nil {
+			log.Println("Error polling folder:", err)
+			continue
+		}
+
+		for _, file := range files {
+			fmt.Println("New .drmd detected via polling:", file)
+			go processWatchedFile(file)
+		}
+	}
+}
+
 func processWatchedFile(filePath string) {
 	for {
 		initialSize, err := getFileSize(filePath)