]> Sergey Matveev's repositories - btrtrc.git/blob - cmd/torrent/main.go
cmd/torrent: Track download rate
[btrtrc.git] / cmd / torrent / main.go
1 // Downloads torrents from the command-line.
2 package main
3
4 import (
5         "expvar"
6         "fmt"
7         "io"
8         stdLog "log"
9         "net"
10         "net/http"
11         "os"
12         "os/signal"
13         "strings"
14         "syscall"
15         "time"
16
17         "github.com/alexflint/go-arg"
18         "github.com/anacrolix/missinggo"
19         "github.com/anacrolix/torrent/bencode"
20         "github.com/davecgh/go-spew/spew"
21         "github.com/dustin/go-humanize"
22         "golang.org/x/xerrors"
23
24         "github.com/anacrolix/log"
25
26         "github.com/anacrolix/envpprof"
27         "github.com/anacrolix/tagflag"
28         "golang.org/x/time/rate"
29
30         "github.com/anacrolix/torrent"
31         "github.com/anacrolix/torrent/iplist"
32         "github.com/anacrolix/torrent/metainfo"
33         "github.com/anacrolix/torrent/storage"
34 )
35
36 func torrentBar(t *torrent.Torrent, pieceStates bool) {
37         go func() {
38                 start := time.Now()
39                 if t.Info() == nil {
40                         fmt.Printf("%v: getting torrent info for %q\n", time.Since(start), t.Name())
41                         <-t.GotInfo()
42                 }
43                 lastStats := t.Stats()
44                 var lastLine string
45                 for range time.Tick(time.Second) {
46                         var completedPieces, partialPieces int
47                         psrs := t.PieceStateRuns()
48                         for _, r := range psrs {
49                                 if r.Complete {
50                                         completedPieces += r.Length
51                                 }
52                                 if r.Partial {
53                                         partialPieces += r.Length
54                                 }
55                         }
56                         stats := t.Stats()
57                         line := fmt.Sprintf(
58                                 "%v: downloading %q: %s/%s, %d/%d pieces completed (%d partial): %v/s\n",
59                                 time.Since(start),
60                                 t.Name(),
61                                 humanize.Bytes(uint64(t.BytesCompleted())),
62                                 humanize.Bytes(uint64(t.Length())),
63                                 completedPieces,
64                                 t.NumPieces(),
65                                 partialPieces,
66                                 humanize.Bytes(uint64(stats.BytesReadUsefulData.Int64()-lastStats.BytesReadUsefulData.Int64())),
67                         )
68                         if line != lastLine {
69                                 lastLine = line
70                                 os.Stdout.WriteString(line)
71                         }
72                         if pieceStates {
73                                 fmt.Println(psrs)
74                         }
75                         lastStats = stats
76                 }
77         }()
78 }
79
80 type stringAddr string
81
82 func (stringAddr) Network() string   { return "" }
83 func (me stringAddr) String() string { return string(me) }
84
85 func resolveTestPeers(addrs []string) (ret []torrent.PeerInfo) {
86         for _, ta := range flags.TestPeer {
87                 ret = append(ret, torrent.PeerInfo{
88                         Addr: stringAddr(ta),
89                 })
90         }
91         return
92 }
93
94 func addTorrents(client *torrent.Client) error {
95         testPeers := resolveTestPeers(flags.TestPeer)
96         for _, arg := range flags.Torrent {
97                 t, err := func() (*torrent.Torrent, error) {
98                         if strings.HasPrefix(arg, "magnet:") {
99                                 t, err := client.AddMagnet(arg)
100                                 if err != nil {
101                                         return nil, xerrors.Errorf("error adding magnet: %w", err)
102                                 }
103                                 return t, nil
104                         } else if strings.HasPrefix(arg, "http://") || strings.HasPrefix(arg, "https://") {
105                                 response, err := http.Get(arg)
106                                 if err != nil {
107                                         return nil, xerrors.Errorf("Error downloading torrent file: %s", err)
108                                 }
109
110                                 metaInfo, err := metainfo.Load(response.Body)
111                                 defer response.Body.Close()
112                                 if err != nil {
113                                         return nil, xerrors.Errorf("error loading torrent file %q: %s\n", arg, err)
114                                 }
115                                 t, err := client.AddTorrent(metaInfo)
116                                 if err != nil {
117                                         return nil, xerrors.Errorf("adding torrent: %w", err)
118                                 }
119                                 return t, nil
120                         } else if strings.HasPrefix(arg, "infohash:") {
121                                 t, _ := client.AddTorrentInfoHash(metainfo.NewHashFromHex(strings.TrimPrefix(arg, "infohash:")))
122                                 return t, nil
123                         } else {
124                                 metaInfo, err := metainfo.LoadFromFile(arg)
125                                 if err != nil {
126                                         return nil, xerrors.Errorf("error loading torrent file %q: %s\n", arg, err)
127                                 }
128                                 t, err := client.AddTorrent(metaInfo)
129                                 if err != nil {
130                                         return nil, xerrors.Errorf("adding torrent: %w", err)
131                                 }
132                                 return t, nil
133                         }
134                 }()
135                 if err != nil {
136                         return xerrors.Errorf("adding torrent for %q: %w", arg, err)
137                 }
138                 if flags.Progress {
139                         torrentBar(t, flags.PieceStates)
140                 }
141                 t.AddPeers(testPeers)
142                 go func() {
143                         <-t.GotInfo()
144                         if len(flags.File) == 0 {
145                                 t.DownloadAll()
146                         } else {
147                                 for _, f := range t.Files() {
148                                         for _, fileArg := range flags.File {
149                                                 if f.DisplayPath() == fileArg {
150                                                         f.Download()
151                                                 }
152                                         }
153                                 }
154                         }
155                 }()
156         }
157         return nil
158 }
159
160 var flags struct {
161         Debug bool
162
163         *DownloadCmd      `arg:"subcommand:download"`
164         *ListFilesCmd     `arg:"subcommand:list-files"`
165         *SpewBencodingCmd `arg:"subcommand:spew-bencoding"`
166         *AnnounceCmd      `arg:"subcommand:announce"`
167 }
168
169 type SpewBencodingCmd struct{}
170
171 type DownloadCmd struct {
172         Mmap            bool           `help:"memory-map torrent data"`
173         TestPeer        []string       `help:"addresses of some starting peers"`
174         Seed            bool           `help:"seed after download is complete"`
175         Addr            string         `help:"network listen addr"`
176         UploadRate      *tagflag.Bytes `help:"max piece bytes to send per second"`
177         DownloadRate    *tagflag.Bytes `help:"max bytes per second down from peers"`
178         PackedBlocklist string
179         PublicIP        net.IP
180         Progress        bool `default:"true"`
181         PieceStates     bool
182         Quiet           bool  `help:"discard client logging"`
183         Stats           *bool `help:"print stats at termination"`
184         Dht             bool  `default:"true"`
185
186         TcpPeers        bool `default:"true"`
187         UtpPeers        bool `default:"true"`
188         Webtorrent      bool `default:"true"`
189         DisableWebseeds bool
190
191         Ipv4 bool `default:"true"`
192         Ipv6 bool `default:"true"`
193         Pex  bool `default:"true"`
194
195         File    []string
196         Torrent []string `arity:"+" help:"torrent file path or magnet uri" arg:"positional"`
197 }
198
199 type ListFilesCmd struct {
200         TorrentPath string `arg:"positional"`
201 }
202
203 func stdoutAndStderrAreSameFile() bool {
204         fi1, _ := os.Stdout.Stat()
205         fi2, _ := os.Stderr.Stat()
206         return os.SameFile(fi1, fi2)
207 }
208
209 func statsEnabled() bool {
210         if flags.Stats == nil {
211                 return flags.Debug
212         }
213         return *flags.Stats
214 }
215
216 func exitSignalHandlers(notify *missinggo.SynchronizedEvent) {
217         c := make(chan os.Signal, 1)
218         signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
219         for {
220                 log.Printf("close signal received: %+v", <-c)
221                 notify.Set()
222         }
223 }
224
225 func main() {
226         if err := mainErr(); err != nil {
227                 log.Printf("error in main: %v", err)
228                 os.Exit(1)
229         }
230 }
231
232 func mainErr() error {
233         stdLog.SetFlags(stdLog.Flags() | stdLog.Lshortfile)
234         p := arg.MustParse(&flags)
235         switch {
236         case flags.AnnounceCmd != nil:
237                 return announceErr()
238         //case :
239         //      return announceErr(flags.Args, parser)
240         case flags.DownloadCmd != nil:
241                 return downloadErr()
242         case flags.ListFilesCmd != nil:
243                 mi, err := metainfo.LoadFromFile(flags.ListFilesCmd.TorrentPath)
244                 if err != nil {
245                         return fmt.Errorf("loading from file %q: %v", flags.ListFilesCmd.TorrentPath, err)
246                 }
247                 info, err := mi.UnmarshalInfo()
248                 if err != nil {
249                         return fmt.Errorf("unmarshalling info from metainfo at %q: %v", flags.ListFilesCmd.TorrentPath, err)
250                 }
251                 for _, f := range info.UpvertedFiles() {
252                         fmt.Println(f.DisplayPath(&info))
253                 }
254                 return nil
255         case flags.SpewBencodingCmd != nil:
256                 d := bencode.NewDecoder(os.Stdin)
257                 for i := 0; ; i++ {
258                         var v interface{}
259                         err := d.Decode(&v)
260                         if err == io.EOF {
261                                 break
262                         }
263                         if err != nil {
264                                 return fmt.Errorf("decoding message index %d: %w", i, err)
265                         }
266                         spew.Dump(v)
267                 }
268                 return nil
269         default:
270                 p.Fail(fmt.Sprintf("unexpected subcommand: %v", p.Subcommand()))
271                 panic("unreachable")
272         }
273 }
274
275 func downloadErr() error {
276         defer envpprof.Stop()
277         clientConfig := torrent.NewDefaultClientConfig()
278         clientConfig.DisableWebseeds = flags.DisableWebseeds
279         clientConfig.DisableTCP = !flags.TcpPeers
280         clientConfig.DisableUTP = !flags.UtpPeers
281         clientConfig.DisableIPv4 = !flags.Ipv4
282         clientConfig.DisableIPv6 = !flags.Ipv6
283         clientConfig.DisableAcceptRateLimiting = true
284         clientConfig.NoDHT = !flags.Dht
285         clientConfig.Debug = flags.Debug
286         clientConfig.Seed = flags.Seed
287         clientConfig.PublicIp4 = flags.PublicIP
288         clientConfig.PublicIp6 = flags.PublicIP
289         clientConfig.DisablePEX = !flags.Pex
290         clientConfig.DisableWebtorrent = !flags.Webtorrent
291         if flags.PackedBlocklist != "" {
292                 blocklist, err := iplist.MMapPackedFile(flags.PackedBlocklist)
293                 if err != nil {
294                         return xerrors.Errorf("loading blocklist: %v", err)
295                 }
296                 defer blocklist.Close()
297                 clientConfig.IPBlocklist = blocklist
298         }
299         if flags.Mmap {
300                 clientConfig.DefaultStorage = storage.NewMMap("")
301         }
302         if flags.Addr != "" {
303                 clientConfig.SetListenAddr(flags.Addr)
304         }
305         if flags.UploadRate != nil {
306                 clientConfig.UploadRateLimiter = rate.NewLimiter(rate.Limit(*flags.UploadRate), 256<<10)
307         }
308         if flags.DownloadRate != nil {
309                 clientConfig.DownloadRateLimiter = rate.NewLimiter(rate.Limit(*flags.DownloadRate), 1<<20)
310         }
311         if flags.Quiet {
312                 clientConfig.Logger = log.Discard
313         }
314
315         var stop missinggo.SynchronizedEvent
316         defer func() {
317                 stop.Set()
318         }()
319
320         client, err := torrent.NewClient(clientConfig)
321         if err != nil {
322                 return xerrors.Errorf("creating client: %v", err)
323         }
324         defer client.Close()
325         go exitSignalHandlers(&stop)
326         go func() {
327                 <-stop.C()
328                 client.Close()
329         }()
330
331         // Write status on the root path on the default HTTP muxer. This will be bound to localhost
332         // somewhere if GOPPROF is set, thanks to the envpprof import.
333         http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
334                 client.WriteStatus(w)
335         })
336         err = addTorrents(client)
337         if err != nil {
338                 return fmt.Errorf("adding torrents: %w", err)
339         }
340         defer outputStats(client)
341         if client.WaitAll() {
342                 log.Print("downloaded ALL the torrents")
343         } else {
344                 return xerrors.New("y u no complete torrents?!")
345         }
346         if flags.Seed {
347                 outputStats(client)
348                 <-stop.C()
349         }
350         return nil
351 }
352
353 func outputStats(cl *torrent.Client) {
354         if !statsEnabled() {
355                 return
356         }
357         expvar.Do(func(kv expvar.KeyValue) {
358                 fmt.Printf("%s: %s\n", kv.Key, kv.Value)
359         })
360         cl.WriteStatus(os.Stdout)
361 }