From b66b5f8e75a21a303a12c0b34737b763d8fad180 Mon Sep 17 00:00:00 2001 From: rumanzo Date: Sat, 9 Apr 2022 00:09:33 +0300 Subject: [PATCH] Rewrite handle trackers function. Respect trackers tier Expect for different types in resumeItem trackers. New recursion function for dig interface for strings in helpers https://github.com/rumanzo/bt2qbt/issues/36 --- bt2qbt.go | 2 +- internal/transfer/fastresumeHandle.go | 2 +- internal/transfer/transfer.go | 38 ++++++++++------- internal/transfer/transfer_test.go | 60 +++++++++++++++++---------- pkg/helpers/helpers.go | 18 ++++++++ pkg/helpers/helpers_test.go | 29 +++++++++++++ pkg/utorrentStructs/utorrent.go | 2 +- 7 files changed, 110 insertions(+), 41 deletions(-) create mode 100644 pkg/helpers/helpers_test.go diff --git a/bt2qbt.go b/bt2qbt.go index 54b5eb0..dd49807 100644 --- a/bt2qbt.go +++ b/bt2qbt.go @@ -46,7 +46,7 @@ func main() { resumeItems := map[string]*utorrentStructs.ResumeItem{} err = bencode.DecodeBytes(b, &resumeItems) if err != nil { - log.Println("Can't process convert resume.dat") + log.Printf("Can't convert resume.dat. Err: %v\n", err) time.Sleep(30 * time.Second) os.Exit(1) } diff --git a/internal/transfer/fastresumeHandle.go b/internal/transfer/fastresumeHandle.go index a88810c..b1c14a0 100644 --- a/internal/transfer/fastresumeHandle.go +++ b/internal/transfer/fastresumeHandle.go @@ -29,7 +29,7 @@ func (transfer *TransferStructure) HandleStructures() { transfer.HandleTags() transfer.HandleLabels() - transfer.GetTrackers(transfer.ResumeItem.Trackers) + transfer.HandleTrackers() transfer.HandlePriority() // important handle priorities before handling pieces /* diff --git a/internal/transfer/transfer.go b/internal/transfer/transfer.go index a6e0f25..39357db 100644 --- a/internal/transfer/transfer.go +++ b/internal/transfer/transfer.go @@ -6,11 +6,13 @@ import ( "github.com/rumanzo/bt2qbt/internal/options" "github.com/rumanzo/bt2qbt/internal/replace" "github.com/rumanzo/bt2qbt/pkg/fileHelpers" + "github.com/rumanzo/bt2qbt/pkg/helpers" "github.com/rumanzo/bt2qbt/pkg/qBittorrentStructures" "github.com/rumanzo/bt2qbt/pkg/torrentStructures" "github.com/rumanzo/bt2qbt/pkg/utorrentStructs" "github.com/zeebo/bencode" "io" + "regexp" "strings" "time" ) @@ -141,24 +143,30 @@ func (transfer *TransferStructure) HandleLabels() { } } -// GetTrackers recurstive function for searching trackers in resume item trackers -func (transfer *TransferStructure) GetTrackers(trackers interface{}) { - switch strct := trackers.(type) { - case []string: - for _, str := range strct { - for _, substr := range strings.Fields(str) { - transfer.Fastresume.Trackers = append(transfer.Fastresume.Trackers, []string{substr}) - } +var localTracker = regexp.MustCompile(`(http|udp)://\S+\.local\S*`) + +func (transfer *TransferStructure) HandleTrackers() { + trackers := helpers.GetStrings(transfer.ResumeItem.Trackers) + trackersMap := map[string][]string{} + var index string + for _, tracker := range trackers { + if localTracker.MatchString(tracker) { + index = "local" + } else { + index = "main" } - case string: - for _, substr := range strings.Fields(strct) { - transfer.Fastresume.Trackers = append(transfer.Fastresume.Trackers, []string{substr}) - } - case []interface{}: - for _, st := range strct { - transfer.GetTrackers(st) + if _, ok := trackersMap[index]; ok { + trackersMap[index] = append(trackersMap[index], tracker) + } else { + trackersMap[index] = []string{tracker} } } + if val, ok := trackersMap["main"]; ok { + transfer.Fastresume.Trackers = append(transfer.Fastresume.Trackers, val) + } + if val, ok := trackersMap["local"]; ok { + transfer.Fastresume.Trackers = append(transfer.Fastresume.Trackers, val) + } } func (transfer *TransferStructure) HandlePriority() { diff --git a/internal/transfer/transfer_test.go b/internal/transfer/transfer_test.go index 24fafb2..3e24c44 100644 --- a/internal/transfer/transfer_test.go +++ b/internal/transfer/transfer_test.go @@ -1124,36 +1124,50 @@ func TestTransferStructure_HandlePriority(t *testing.T) { } } -func TestTransferStructure_GetTrackers(t *testing.T) { +func TestTransferStructure_HandleTrackers(t *testing.T) { transferStructure := TransferStructure{ Fastresume: &qBittorrentStructures.QBittorrentFastresume{}, - } - testTrackers := []interface{}{ - "test1", - "test2", - []interface{}{ - "test3", - "test4", - []interface{}{ - "test5", - "test6", + ResumeItem: &utorrentStructs.ResumeItem{ + Trackers: []interface{}{ + "http://test1.org", + "udp://test1.org", + "http://test1.local", + "udp://test1.local", + []interface{}{ + "http://test2.org:80", + "udp://test2.org:8080", + "http://test2.local:80", + "udp://test2.local:8080", + []interface{}{ + "http://test3.org:80/somepath", + "udp://test3.org:8080/somepath", + "http://test3.local:80/somepath", + "udp://test3.local:8080/somepath", + }, + }, + []interface{}{ + []interface{}{ + "http://test4.org:80/", + "udp://test4.org:8080/", + "http://test4.local:80/", + "udp://test4.local:8080/", + }, + }, }, }, - []interface{}{ - []interface{}{"test7", "test8"}, - }, } expect := [][]string{ - []string{"test1"}, - []string{"test2"}, - []string{"test3"}, - []string{"test4"}, - []string{"test5"}, - []string{"test6"}, - []string{"test7"}, - []string{"test8"}, + []string{ + "http://test1.org", "udp://test1.org", + "http://test2.org:80", "udp://test2.org:8080", + "http://test3.org:80/somepath", "udp://test3.org:8080/somepath", + "http://test4.org:80/", "udp://test4.org:8080/"}, + []string{"http://test1.local", "udp://test1.local", + "http://test2.local:80", "udp://test2.local:8080", + "http://test3.local:80/somepath", "udp://test3.local:8080/somepath", + "http://test4.local:80/", "udp://test4.local:8080/"}, } - transferStructure.GetTrackers(testTrackers) + transferStructure.HandleTrackers() if !reflect.DeepEqual(transferStructure.Fastresume.Trackers, expect) { t.Fatalf("Unexpected error: opts isn't equal:\n Got: %#v\n Expect %#v\n", transferStructure.Fastresume.Trackers, expect) } diff --git a/pkg/helpers/helpers.go b/pkg/helpers/helpers.go index 50d1271..87d47bd 100644 --- a/pkg/helpers/helpers.go +++ b/pkg/helpers/helpers.go @@ -8,6 +8,7 @@ import ( "io/ioutil" "os" "strconv" + "strings" ) func ASCIIConvert(s string) string { @@ -93,3 +94,20 @@ func CopyFile(src string, dst string) error { } return nil } + +func GetStrings(trackers interface{}) []string { + ntrackers := []string{} + switch strct := trackers.(type) { + case []string: + for _, str := range strct { + ntrackers = append(ntrackers, strings.Fields(str)...) + } + case string: + ntrackers = append(ntrackers, strings.Fields(strct)...) + case []interface{}: + for _, st := range strct { + ntrackers = append(ntrackers, GetStrings(st)...) + } + } + return ntrackers +} diff --git a/pkg/helpers/helpers_test.go b/pkg/helpers/helpers_test.go new file mode 100644 index 0000000..b8027e9 --- /dev/null +++ b/pkg/helpers/helpers_test.go @@ -0,0 +1,29 @@ +package helpers + +import ( + "reflect" + "testing" +) + +func TestGetStrings(t *testing.T) { + testTrackers := []interface{}{ + "test1", + "test2", + []interface{}{ + "test3", + "test4", + []interface{}{ + "test5", + "test6", + }, + }, + []interface{}{ + []interface{}{"test7", "test8"}, + }, + } + expect := []string{"test1", "test2", "test3", "test4", "test5", "test6", "test7", "test8"} + trackers := GetStrings(testTrackers) + if !reflect.DeepEqual(trackers, expect) { + t.Fatalf("Unexpected error: opts isn't equal:\n Got: %#v\n Expect %#v\n", trackers, expect) + } +} diff --git a/pkg/utorrentStructs/utorrent.go b/pkg/utorrentStructs/utorrent.go index 3bc549e..ed87121 100644 --- a/pkg/utorrentStructs/utorrent.go +++ b/pkg/utorrentStructs/utorrent.go @@ -15,7 +15,7 @@ type ResumeItem struct { Started int64 `bencode:"started"` Targets [][]interface{} `bencode:"targets,omitempty"` Time int64 `bencode:"time"` - Trackers []string `bencode:"trackers,omitempty"` + Trackers interface{} `bencode:"trackers,omitempty"` UpSpeed int64 `bencode:"upspeed"` Uploaded int64 `bencode:"uploaded"` }