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