Refactor structure, use cobra

cmd/zurg/main.go

@@ -2,79 +2,45 @@ package main
 
 import (
 	"fmt"
-	"net/http"
 	"os"
 
-	"github.com/debridmediamanager.com/zurg/internal/config"
+	"github.com/debridmediamanager.com/zurg/internal"
-	"github.com/debridmediamanager.com/zurg/internal/net"
+	"github.com/spf13/cobra"
-	"github.com/debridmediamanager.com/zurg/internal/torrent"
-	"github.com/debridmediamanager.com/zurg/internal/universal"
-	"github.com/debridmediamanager.com/zurg/internal/version"
-	"github.com/dgraph-io/ristretto"
-	"github.com/panjf2000/ants/v2"
-
-	zurghttp "github.com/debridmediamanager.com/zurg/pkg/http"
-	"github.com/debridmediamanager.com/zurg/pkg/logutil"
-	"github.com/debridmediamanager.com/zurg/pkg/realdebrid"
 )
 
 func main() {
-	// special commands
+	var configPath string // Variable to hold the config path
-	if len(os.Args) > 1 {
+
-		switch os.Args[1] {
+	var rootCmd = &cobra.Command{
-		case "version":
+		Use: "zurg",
-			version.Show()
+		Run: func(cmd *cobra.Command, args []string) {
-		case "networktest":
+			internal.MainApp(configPath) // Pass the configPath to MainApp
-			realdebrid.RunTest()
+		},
-		}
-		os.Exit(0)
 	}
 
-	// normal startup
+	// Adding a flag for the config path with a default value
-	log := logutil.NewLogger()
+	rootCmd.PersistentFlags().StringVarP(&configPath, "config", "c", "./config.yml", "config file path")
-	zurglog := log.Named("zurg")
 
-	config, configErr := config.LoadZurgConfig("./config.yml", log.Named("config"))
+	var versionCmd = &cobra.Command{
-	if configErr != nil {
+		Use:   "version",
-		zurglog.Errorf("Config failed to load: %v", configErr)
+		Short: "Prints zurg's current version",
-		os.Exit(1)
+		Run: func(cmd *cobra.Command, args []string) {
+			internal.ShowVersion()
+		},
 	}
 
-	apiClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), config.GetRealDebridTimeout(), config, log.Named("httpclient"))
+	var networkTestCmd = &cobra.Command{
-
+		Use:   "network-test",
-	rd := realdebrid.NewRealDebrid(apiClient, log.Named("realdebrid"))
+		Short: "Run a network test",
-
+		Run: func(cmd *cobra.Command, args []string) {
-	p, err := ants.NewPool(config.GetNumOfWorkers())
+			internal.NetworkTest()
-	if err != nil {
+		},
-		zurglog.Errorf("Failed to create worker pool: %v", err)
-		os.Exit(1)
-	}
-	defer p.Release()
-
-	cache, err := ristretto.NewCache(&ristretto.Config{
-		NumCounters: 1e5,       // 200,000 to track frequency for 100,000 items.
-		MaxCost:     100 << 20, // maximum cost of cache (100MB).
-		BufferItems: 1 << 10,   // number of keys per Get buffer, can be adjusted.
-	})
-	if err != nil {
-		zurglog.Errorf("Failed to create cache: %v", err)
-		os.Exit(1)
 	}
 
-	torrentMgr := torrent.NewTorrentManager(config, rd, p, cache, log.Named("manager"))
+	rootCmd.AddCommand(versionCmd, networkTestCmd)
 
-	downloadClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), 0, config, log.Named("dlclient"))
+	if err := rootCmd.Execute(); err != nil {
-	getfile := universal.NewGetFile(downloadClient)
+		fmt.Println(err)
-
-	mux := http.NewServeMux()
-	net.Router(mux, getfile, config, torrentMgr, log.Named("net"))
-
-	addr := fmt.Sprintf("%s:%s", config.GetHost(), config.GetPort())
-	server := &http.Server{Addr: addr, Handler: mux}
-
-	zurglog.Infof("Starting server on %s", addr)
-	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
-		zurglog.Errorf("Failed to start server: %v", err)
 		os.Exit(1)
 	}
 }

go.mod

@@ -13,10 +13,13 @@ require github.com/panjf2000/ants/v2 v2.8.2
 
 require (
 	github.com/cespare/xxhash/v2 v2.1.1 // indirect
-	github.com/dgraph-io/ristretto v0.1.1 // indirect
+	github.com/dgraph-io/ristretto v0.1.1
 	github.com/dustin/go-humanize v1.0.0 // indirect
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b // indirect
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
+	github.com/spf13/cobra v1.8.0
+	github.com/spf13/pflag v1.0.5 // indirect
 	go.uber.org/multierr v1.10.0 // indirect
 	golang.org/x/sys v0.0.0-20221010170243-090e33056c14 // indirect
 )

go.sum

@@ -1,15 +1,19 @@
 github.com/cespare/xxhash/v2 v2.1.1 h1:6MnRN8NT7+YBpUIWxHtefFZOKTAPgGjpQSxqLNn0+qY=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgraph-io/ristretto v0.1.1 h1:6CWw5tJNgpegArSHpNHJKldNeq03FQCwYvfMVWajOK8=
 github.com/dgraph-io/ristretto v0.1.1/go.mod h1:S1GPSBCYCIhmVNfcth17y2zZtQT6wzkzgwUve0VDWWA=
+github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 h1:tdlZCpZ/P9DhczCTSixgIKmwPv6+wP5DGjqLYw5SUiA=
 github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/orcaman/concurrent-map/v2 v2.0.1 h1:jOJ5Pg2w1oeB6PeDurIYf6k9PQ+aTITr/6lP/L/zp6c=
 github.com/orcaman/concurrent-map/v2 v2.0.1/go.mod h1:9Eq3TG2oBe5FirmYWQfYO5iH1q0Jv47PLaNK++uCdOM=
 github.com/panjf2000/ants/v2 v2.8.2 h1:D1wfANttg8uXhC9149gRt1PDQ+dLVFjNXkCEycMcvQQ=
@@ -18,6 +22,11 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0=
+github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=

internal/app.go

@@ -0,0 +1,69 @@
+package internal
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+
+	"github.com/debridmediamanager.com/zurg/internal/config"
+	"github.com/debridmediamanager.com/zurg/internal/net"
+	"github.com/debridmediamanager.com/zurg/internal/torrent"
+	"github.com/debridmediamanager.com/zurg/internal/universal"
+	zurghttp "github.com/debridmediamanager.com/zurg/pkg/http"
+	"github.com/debridmediamanager.com/zurg/pkg/logutil"
+	"github.com/debridmediamanager.com/zurg/pkg/realdebrid"
+	"github.com/debridmediamanager.com/zurg/pkg/utils"
+	"github.com/dgraph-io/ristretto"
+	"github.com/panjf2000/ants/v2"
+)
+
+func MainApp(configPath string) {
+	log := logutil.NewLogger()
+	zurglog := log.Named("zurg")
+
+	config, configErr := config.LoadZurgConfig(configPath, log.Named("config"))
+	if configErr != nil {
+		zurglog.Errorf("Config failed to load: %v", configErr)
+		os.Exit(1)
+	}
+
+	apiClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), config.GetRealDebridTimeout(), config, log.Named("httpclient"))
+
+	rd := realdebrid.NewRealDebrid(apiClient, log.Named("realdebrid"))
+
+	p, err := ants.NewPool(config.GetNumOfWorkers())
+	if err != nil {
+		zurglog.Errorf("Failed to create worker pool: %v", err)
+		os.Exit(1)
+	}
+	defer p.Release()
+
+	utils.EnsureDirExists("data")
+
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1e5,       // 200,000 to track frequency for 100,000 items.
+		MaxCost:     100 << 20, // maximum cost of cache (100MB).
+		BufferItems: 1 << 10,   // number of keys per Get buffer, can be adjusted.
+	})
+	if err != nil {
+		zurglog.Errorf("Failed to create cache: %v", err)
+		os.Exit(1)
+	}
+
+	torrentMgr := torrent.NewTorrentManager(config, rd, p, cache, log.Named("manager"))
+
+	downloadClient := zurghttp.NewHTTPClient(config.GetToken(), config.GetRetriesUntilFailed(), 0, config, log.Named("dlclient"))
+	getfile := universal.NewGetFile(downloadClient)
+
+	mux := http.NewServeMux()
+	net.Router(mux, getfile, config, torrentMgr, log.Named("net"))
+
+	addr := fmt.Sprintf("%s:%s", config.GetHost(), config.GetPort())
+	server := &http.Server{Addr: addr, Handler: mux}
+
+	zurglog.Infof("Starting server on %s", addr)
+	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		zurglog.Errorf("Failed to start server: %v", err)
+		os.Exit(1)
+	}
+}

internal/network.go

@@ -0,0 +1,7 @@
+package internal
+
+import "github.com/debridmediamanager.com/zurg/pkg/realdebrid"
+
+func NetworkTest() {
+	realdebrid.RunTest()
+}

internal/torrent/manager.go

@@ -40,7 +40,6 @@ type TorrentManager struct {
 	antsPool        *ants.Pool
 	unrestrictPool  *ants.Pool
 	log             *zap.SugaredLogger
-	mu              *sync.Mutex
 }
 
 // NewTorrentManager creates a new torrent manager
@@ -55,7 +54,6 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
 		requiredVersion: "18.11.2023",
 		antsPool:        p,
 		log:             log,
-		mu:              &sync.Mutex{},
 	}
 
 	unrestrictPool, err := ants.NewPool(t.Config.GetUnrestrictWorkers())
@@ -65,8 +63,6 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
 		t.unrestrictPool = unrestrictPool
 	}
 
-	ensureDir(DATA_DIR)
-
 	// create internal directories
 	t.DirectoryMap.Set(INT_ALL, cmap.New[*Torrent]())        // key is AccessKey
 	t.DirectoryMap.Set(INT_INFO_CACHE, cmap.New[*Torrent]()) // key is Torrent ID

internal/torrent/util.go

@@ -2,7 +2,6 @@ package torrent
 
 import (
 	"fmt"
-	"os"
 )
 
 func getFileIDs(files []File) []string {
@@ -14,15 +13,3 @@ func getFileIDs(files []File) []string {
 	}
 	return fileIDs
 }
-
-func ensureDir(dirName string) error {
-	if _, err := os.Stat(dirName); os.IsNotExist(err) {
-		err := os.Mkdir(dirName, 0755)
-		if err != nil {
-			return err
-		}
-	} else if err != nil {
-		return err
-	}
-	return nil
-}

internal/version.go

@@ -1,4 +1,4 @@
-package version
+package internal
 
 import (
 	"fmt"
@@ -11,7 +11,7 @@ var (
 	Version   string = "dev"
 )
 
-func Show() {
+func ShowVersion() {
 	fmt.Printf("zurg\nBuilt At: %s\nGo Version: %s\nCommit: %s\nVersion: %s\n",
 		BuiltAt, GoVersion, GitCommit, Version)
 }

pkg/realdebrid/network.go

@@ -31,7 +31,7 @@ func traceroute(ip string) (int, time.Duration, error) {
 	var latency time.Duration
 	if hopCount > 0 {
 		lastLine := lines[hopCount-1]
-		fmt.Println(lastLine)
+		fmt.Printf("IP: %s, last hop: %s\n", ip, lastLine)
 		parts := strings.Fields(lastLine)
 		if len(parts) >= 3 {
 			latencyValue, parseErr := strconv.ParseFloat(parts[2], 64)
@@ -70,6 +70,7 @@ func RunTest() {
 
 	wg.Wait()
 	close(infoChan)
+	fmt.Printf("Network test complete.\n\n")
 
 	var ipInfos []IPInfo
 	for info := range infoChan {
@@ -77,30 +78,36 @@ func RunTest() {
 	}
 
 	sort.Slice(ipInfos, func(i, j int) bool {
-		if ipInfos[i].Hops == ipInfos[j].Hops {
 		return ipInfos[i].Latency < ipInfos[j].Latency
-		}
-		return ipInfos[i].Hops < ipInfos[j].Hops
 	})
-	var lowestLatency time.Duration
-	if len(ipInfos) > 0 {
-		lowestLatency = ipInfos[0].Latency
-		for _, info := range ipInfos {
-			if info.Latency < lowestLatency {
-				lowestLatency = info.Latency
-			}
-		}
-	}
-
-	latencyThreshold := lowestLatency + lowestLatency/3
 
+	const minResults = 10
+	const maxResults = 20
 	var okIPs []IPInfo
-	for _, info := range ipInfos {
+
-		if info.Latency <= latencyThreshold {
+	if len(ipInfos) > 0 {
+		// Start by adding the best IPs based on hops and latency up to the minResults
+		for i := 0; i < min(len(ipInfos), minResults); i++ {
+			okIPs = append(okIPs, ipInfos[i])
+		}
+
+		// Find the highest latency in the current okIPs list
+		highestLatency := okIPs[len(okIPs)-1].Latency
+
+		// Add any additional IPs that have latency within a reasonable range of the highest latency
+		for _, info := range ipInfos[minResults:] {
+			if len(okIPs) >= maxResults {
+				break // Stop adding IPs if maxResults is reached
+			}
+			if info.Latency <= highestLatency+(highestLatency/3) {
 				okIPs = append(okIPs, info)
 			}
 		}
+	}
+
+	fmt.Printf("Here are the results, you can copy-paste the following to your config.yml:\n\n")
+	fmt.Println("preferred_hosts:")
 	for _, info := range okIPs {
-		fmt.Printf("Host: %s, Hops: %d, Latency: %v\n", info.Address, info.Hops, info.Latency)
+		fmt.Printf("  - %s # hops: %d latency: %v\n", info.Address, info.Hops, info.Latency)
 	}
 }

pkg/utils/ensure.go

@@ -0,0 +1,15 @@
+package utils
+
+import "os"
+
+func EnsureDirExists(dirName string) error {
+	if _, err := os.Stat(dirName); os.IsNotExist(err) {
+		err := os.Mkdir(dirName, 0755)
+		if err != nil {
+			return err
+		}
+	} else if err != nil {
+		return err
+	}
+	return nil
+}