...
1 package manager
2
3 import (
4 "errors"
5 "strings"
6 "sync"
7
8 "gitlab.com/tslocum/gophast/pkg/config"
9 "gitlab.com/tslocum/gophast/pkg/download"
10 "gitlab.com/tslocum/gophast/pkg/log"
11 )
12
13 type Metadata struct {
14 Name string
15 Size int64
16 }
17
18 type NewDownloadInfo struct {
19 URL string
20 PostData []byte
21 }
22
23 type DownloadInfo struct {
24 Download *download.Download
25 }
26
27 type DownloadResult struct {
28 Status int
29 Err error
30 }
31
32 type ManagedDownload struct {
33 d *download.Download
34 result *DownloadResult
35 wg *sync.WaitGroup
36 }
37
38 func (md *ManagedDownload) download() {
39 defer func(md *ManagedDownload) {
40 md.wg.Done()
41 config.WG.Done()
42 }(md)
43
44 md.d.RLock()
45 if md.d.Cancelled || md.d.Status > 0 {
46 md.d.RUnlock()
47 return
48 }
49 md.d.RUnlock()
50
51 err := md.d.Download(nil)
52
53 if err != nil {
54 if !strings.Contains(err.Error(), "already downloaded") {
55 log.Verbosef("failed to download %s: %v", md.d.Name, err)
56 } else {
57 log.Verbosef("%s %v", md.d.Name, err)
58 }
59 }
60
61 managerLock.Lock()
62 md.result = &DownloadResult{md.d.GetStatus(), err}
63 managerLock.Unlock()
64 }
65
66 var (
67 downloads map[int64]*ManagedDownload
68 managerLock = &sync.RWMutex{}
69 )
70
71 func FetchMetadata(info *NewDownloadInfo) (*Metadata, error) {
72 metadata, err := download.FetchMetadata(info.URL, info.PostData)
73 if err != nil {
74 return nil, err
75 }
76
77 return &Metadata{Name: metadata.Name, Size: metadata.Size}, nil
78 }
79
80 func AddDownload(info *NewDownloadInfo) (int64, error) {
81 if downloads == nil {
82 downloads = make(map[int64]*ManagedDownload)
83 }
84
85 dl, err := download.NewDownload(info.URL, info.PostData)
86 if err != nil {
87 return -1, err
88 }
89
90 log.Verbosef("Initiating download #%d %s", dl.ID, info.URL)
91
92 config.WG.Add(1)
93
94 md := &ManagedDownload{d: dl, wg: &sync.WaitGroup{}}
95
96 managerLock.Lock()
97 downloads[dl.ID] = md
98 managerLock.Unlock()
99
100 md.wg.Add(1)
101 go md.download()
102
103 return dl.ID, nil
104 }
105
106 func GetDownload(id int64) *download.Download {
107 managerLock.RLock()
108 defer managerLock.RUnlock()
109
110 if md, ok := downloads[id]; ok {
111 return md.d
112 }
113
114 return nil
115 }
116
117 func GetName(id int64) string {
118 managerLock.RLock()
119 defer managerLock.RUnlock()
120
121 if md, ok := downloads[id]; ok {
122 return md.d.Name
123 }
124
125 return "invalid"
126 }
127
128 func GetSize(id int64) int64 {
129 managerLock.RLock()
130 defer managerLock.RUnlock()
131
132 if md, ok := downloads[id]; ok {
133 return md.d.Size
134 }
135
136 return -1
137 }
138
139 func GetDownloaded(id int64) int64 {
140 managerLock.RLock()
141 defer managerLock.RUnlock()
142
143 if md, ok := downloads[id]; ok {
144 return md.d.GetDownloaded()
145 }
146
147 return -1
148 }
149
150 func GetVerboseDownloaded(id int64) []byte {
151 managerLock.RLock()
152 defer managerLock.RUnlock()
153
154 if md, ok := downloads[id]; ok {
155 return md.d.GetVerboseDownloaded()
156 }
157
158 return nil
159 }
160
161 func PauseDownload(id int64) error {
162 managerLock.Lock()
163 defer managerLock.Unlock()
164
165 if md, ok := downloads[id]; ok {
166 md.d.Pause()
167 log.Verbosef("Paused #%d", id)
168
169 return nil
170 }
171
172 return errors.New("download not found")
173 }
174
175 func ResumeDownload(id int64) error {
176 managerLock.Lock()
177 if md, ok := downloads[id]; ok {
178 if md.d.Cancelled || md.d.Status > 0 {
179 return nil
180 }
181 managerLock.Unlock()
182
183 if md.d.GetStatus() > 0 {
184 return errors.New("download finished")
185 } else if md.d.Cancelled {
186 return errors.New("download cancelled")
187 } else if md.d.GetStatus() == 0 && !md.d.Paused {
188 return errors.New("download in progress")
189 }
190
191 log.Verbosef("Resuming #%d...", id)
192
193 md.d.Resume()
194
195 return nil
196 }
197 managerLock.Unlock()
198
199 return errors.New("download not found")
200 }
201
202 func WaitDownload(id int64) (int, error) {
203 managerLock.RLock()
204 if md, ok := downloads[id]; ok {
205 managerLock.RUnlock()
206
207 md.wg.Wait()
208
209 managerLock.RLock()
210 defer managerLock.RUnlock()
211
212 if md.result == nil {
213 return -1, errors.New("download corrupt")
214 }
215
216 return md.result.Status, md.result.Err
217 }
218 managerLock.RUnlock()
219
220 return -7, errors.New("cancelled")
221 }
222
223 func RemoveDownload(id int64) error {
224 managerLock.Lock()
225 defer managerLock.Unlock()
226
227 log.Verbosef("Removing download #%d", id)
228 if md, ok := downloads[id]; ok {
229 md.d.Cancel()
230
231 delete(downloads, id)
232 return nil
233 }
234
235 return errors.New("failed to remove download: not found")
236 }
237
View as plain text