cmd/torrent/main.go

   1 // Downloads torrents from the command-line.
   2 package main
   3
   4 import (
   5         "expvar"
   6         "fmt"
   7         "log"
   8         "net"
   9         "net/http"
  10         "os"
  11         "os/signal"
  12         "strings"
  13         "syscall"
  14         "time"
  15
  16         "github.com/anacrolix/envpprof"
  17         "github.com/anacrolix/tagflag"
  18         "github.com/anacrolix/torrent"
  19         "github.com/anacrolix/torrent/iplist"
  20         "github.com/anacrolix/torrent/metainfo"
  21         "github.com/anacrolix/torrent/storage"
  22         "github.com/dustin/go-humanize"
  23         "github.com/gosuri/uiprogress"
  24         "golang.org/x/time/rate"
  25 )
  26
  27 var progress = uiprogress.New()
  28
  29 func torrentBar(t *torrent.Torrent) {
  30         bar := progress.AddBar(1)
  31         bar.AppendCompleted()
  32         bar.AppendFunc(func(*uiprogress.Bar) (ret string) {
  33                 select {
  34                 case <-t.GotInfo():
  35                 default:
  36                         return "getting info"
  37                 }
  38                 if t.Seeding() {
  39                         return "seeding"
  40                 } else if t.BytesCompleted() == t.Info().TotalLength() {
  41                         return "completed"
  42                 } else {
  43                         return fmt.Sprintf("downloading (%s/%s)", humanize.Bytes(uint64(t.BytesCompleted())), humanize.Bytes(uint64(t.Info().TotalLength())))
  44                 }
  45         })
  46         bar.PrependFunc(func(*uiprogress.Bar) string {
  47                 return t.Name()
  48         })
  49         go func() {
  50                 <-t.GotInfo()
  51                 tl := int(t.Info().TotalLength())
  52                 if tl == 0 {
  53                         bar.Set(1)
  54                         return
  55                 }
  56                 bar.Total = tl
  57                 for {
  58                         bc := t.BytesCompleted()
  59                         bar.Set(int(bc))
  60                         time.Sleep(time.Second)
  61                 }
  62         }()
  63 }
  64
  65 func addTorrents(client *torrent.Client) {
  66         for _, arg := range flags.Torrent {
  67                 t := func() *torrent.Torrent {
  68                         if strings.HasPrefix(arg, "magnet:") {
  69                                 t, err := client.AddMagnet(arg)
  70                                 if err != nil {
  71                                         log.Fatalf("error adding magnet: %s", err)
  72                                 }
  73                                 return t
  74                         } else if strings.HasPrefix(arg, "http://") || strings.HasPrefix(arg, "https://") {
  75                                 response, err := http.Get(arg)
  76                                 if err != nil {
  77                                         log.Fatalf("Error downloading torrent file: %s", err)
  78                                 }
  79
  80                                 metaInfo, err := metainfo.Load(response.Body)
  81                                 defer response.Body.Close()
  82                                 if err != nil {
  83                                         fmt.Fprintf(os.Stderr, "error loading torrent file %q: %s\n", arg, err)
  84                                         os.Exit(1)
  85                                 }
  86                                 t, err := client.AddTorrent(metaInfo)
  87                                 if err != nil {
  88                                         log.Fatal(err)
  89                                 }
  90                                 return t
  91                         } else if strings.HasPrefix(arg, "infohash:") {
  92                                 t, _ := client.AddTorrentInfoHash(metainfo.NewHashFromHex(strings.TrimPrefix(arg, "infohash:")))
  93                                 return t
  94                         } else {
  95                                 metaInfo, err := metainfo.LoadFromFile(arg)
  96                                 if err != nil {
  97                                         fmt.Fprintf(os.Stderr, "error loading torrent file %q: %s\n", arg, err)
  98                                         os.Exit(1)
  99                                 }
 100                                 t, err := client.AddTorrent(metaInfo)
 101                                 if err != nil {
 102                                         log.Fatal(err)
 103                                 }
 104                                 return t
 105                         }
 106                 }()
 107                 torrentBar(t)
 108                 t.AddPeers(func() (ret []torrent.Peer) {
 109                         for _, ta := range flags.TestPeer {
 110                                 ret = append(ret, torrent.Peer{
 111                                         IP:   ta.IP,
 112                                         Port: ta.Port,
 113                                 })
 114                         }
 115                         return
 116                 }())
 117                 go func() {
 118                         <-t.GotInfo()
 119                         t.DownloadAll()
 120                 }()
 121         }
 122 }
 123
 124 var flags = struct {
 125         Mmap            bool           `help:"memory-map torrent data"`
 126         TestPeer        []*net.TCPAddr `help:"addresses of some starting peers"`
 127         Seed            bool           `help:"seed after download is complete"`
 128         Addr            *net.TCPAddr   `help:"network listen addr"`
 129         UploadRate      tagflag.Bytes  `help:"max piece bytes to send per second"`
 130         DownloadRate    tagflag.Bytes  `help:"max bytes per second down from peers"`
 131         Debug           bool
 132         PackedBlocklist string
 133         Stats           *bool
 134         tagflag.StartPos
 135         Torrent []string `arity:"+" help:"torrent file path or magnet uri"`
 136 }{
 137         UploadRate:   -1,
 138         DownloadRate: -1,
 139 }
 140
 141 func stdoutAndStderrAreSameFile() bool {
 142         fi1, _ := os.Stdout.Stat()
 143         fi2, _ := os.Stderr.Stat()
 144         return os.SameFile(fi1, fi2)
 145 }
 146
 147 func statsEnabled() bool {
 148         if flags.Stats == nil {
 149                 return flags.Debug
 150         }
 151         return *flags.Stats
 152 }
 153
 154 func exitSignalHandlers(client *torrent.Client) {
 155         c := make(chan os.Signal, 1)
 156         signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
 157         for {
 158                 log.Printf("close signal received: %+v", <-c)
 159                 client.Close()
 160         }
 161 }
 162
 163 func main() {
 164         log.SetFlags(log.LstdFlags | log.Lshortfile)
 165         tagflag.Parse(&flags)
 166         defer envpprof.Stop()
 167         clientConfig := torrent.NewDefaultClientConfig()
 168         clientConfig.Debug = flags.Debug
 169         clientConfig.Seed = flags.Seed
 170         if flags.PackedBlocklist != "" {
 171                 blocklist, err := iplist.MMapPackedFile(flags.PackedBlocklist)
 172                 if err != nil {
 173                         log.Fatalf("error loading blocklist: %s", err)
 174                 }
 175                 defer blocklist.Close()
 176                 clientConfig.IPBlocklist = blocklist
 177         }
 178         if flags.Mmap {
 179                 clientConfig.DefaultStorage = storage.NewMMap("")
 180         }
 181         if flags.Addr != nil {
 182                 clientConfig.SetListenAddr(flags.Addr.String())
 183         }
 184         if flags.UploadRate != -1 {
 185                 clientConfig.UploadRateLimiter = rate.NewLimiter(rate.Limit(flags.UploadRate), 256<<10)
 186         }
 187         if flags.DownloadRate != -1 {
 188                 clientConfig.DownloadRateLimiter = rate.NewLimiter(rate.Limit(flags.DownloadRate), 1<<20)
 189         }
 190
 191         client, err := torrent.NewClient(clientConfig)
 192         if err != nil {
 193                 log.Fatalf("error creating client: %s", err)
 194         }
 195         defer client.Close()
 196         go exitSignalHandlers(client)
 197
 198         // Write status on the root path on the default HTTP muxer. This will be
 199         // bound to localhost somewhere if GOPPROF is set, thanks to the envpprof
 200         // import.
 201         http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
 202                 client.WriteStatus(w)
 203         })
 204         if stdoutAndStderrAreSameFile() {
 205                 log.SetOutput(progress.Bypass())
 206         }
 207         progress.Start()
 208         addTorrents(client)
 209         if client.WaitAll() {
 210                 log.Print("downloaded ALL the torrents")
 211         } else {
 212                 log.Fatal("y u no complete torrents?!")
 213         }
 214         if flags.Seed {
 215                 outputStats(client)
 216                 select {}
 217         }
 218         outputStats(client)
 219 }
 220
 221 func outputStats(cl *torrent.Client) {
 222         if !statsEnabled() {
 223                 return
 224         }
 225         expvar.Do(func(kv expvar.KeyValue) {
 226                 fmt.Printf("%s: %s\n", kv.Key, kv.Value)
 227         })
 228         cl.WriteStatus(os.Stdout)
 229 }