cmd/dht-get-peers/main.go

   1 package main
   2
   3 import (
   4         "flag"
   5         "fmt"
   6         "io"
   7         "log"
   8         "net"
   9         "os"
  10         "os/signal"
  11
  12         _ "github.com/anacrolix/envpprof"
  13
  14         "github.com/anacrolix/torrent/dht"
  15         "github.com/anacrolix/torrent/util"
  16         _ "github.com/anacrolix/torrent/util/profile"
  17 )
  18
  19 var (
  20         tableFileName = flag.String("tableFile", "", "name of file for storing node info")
  21         serveAddr     = flag.String("serveAddr", ":0", "local UDP address")
  22         infoHash      = flag.String("infoHash", "", "torrent infohash")
  23         once          = flag.Bool("once", false, "only do one scrape iteration")
  24
  25         s        *dht.Server
  26         quitting = make(chan struct{})
  27 )
  28
  29 func loadTable() error {
  30         if *tableFileName == "" {
  31                 return nil
  32         }
  33         f, err := os.Open(*tableFileName)
  34         if os.IsNotExist(err) {
  35                 return nil
  36         }
  37         if err != nil {
  38                 return fmt.Errorf("error opening table file: %s", err)
  39         }
  40         defer f.Close()
  41         added := 0
  42         for {
  43                 b := make([]byte, dht.CompactNodeInfoLen)
  44                 _, err := io.ReadFull(f, b)
  45                 if err == io.EOF {
  46                         break
  47                 }
  48                 if err != nil {
  49                         return fmt.Errorf("error reading table file: %s", err)
  50                 }
  51                 var ni dht.NodeInfo
  52                 err = ni.UnmarshalCompact(b)
  53                 if err != nil {
  54                         return fmt.Errorf("error unmarshaling compact node info: %s", err)
  55                 }
  56                 s.AddNode(ni)
  57                 added++
  58         }
  59         log.Printf("loaded %d nodes from table file", added)
  60         return nil
  61 }
  62
  63 func init() {
  64         log.SetFlags(log.LstdFlags | log.Lshortfile)
  65         flag.Parse()
  66         switch len(*infoHash) {
  67         case 20:
  68         case 40:
  69                 _, err := fmt.Sscanf(*infoHash, "%x", infoHash)
  70                 if err != nil {
  71                         log.Fatal(err)
  72                 }
  73         default:
  74                 log.Fatal("require 20 byte infohash")
  75         }
  76         var err error
  77         s, err = dht.NewServer(&dht.ServerConfig{
  78                 Addr: *serveAddr,
  79         })
  80         if err != nil {
  81                 log.Fatal(err)
  82         }
  83         err = loadTable()
  84         if err != nil {
  85                 log.Fatalf("error loading table: %s", err)
  86         }
  87         log.Printf("dht server on %s, ID is %x", s.Addr(), s.ID())
  88         setupSignals()
  89 }
  90
  91 func saveTable() error {
  92         goodNodes := s.Nodes()
  93         if *tableFileName == "" {
  94                 if len(goodNodes) != 0 {
  95                         log.Print("good nodes were discarded because you didn't specify a table file")
  96                 }
  97                 return nil
  98         }
  99         f, err := os.OpenFile(*tableFileName, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
 100         if err != nil {
 101                 return fmt.Errorf("error opening table file: %s", err)
 102         }
 103         defer f.Close()
 104         for _, nodeInfo := range goodNodes {
 105                 var b [dht.CompactNodeInfoLen]byte
 106                 err := nodeInfo.PutCompact(b[:])
 107                 if err != nil {
 108                         return fmt.Errorf("error compacting node info: %s", err)
 109                 }
 110                 _, err = f.Write(b[:])
 111                 if err != nil {
 112                         return fmt.Errorf("error writing compact node info: %s", err)
 113                 }
 114         }
 115         log.Printf("saved %d nodes to table file", len(goodNodes))
 116         return nil
 117 }
 118
 119 func setupSignals() {
 120         ch := make(chan os.Signal)
 121         signal.Notify(ch, os.Interrupt)
 122         go func() {
 123                 <-ch
 124                 close(quitting)
 125         }()
 126 }
 127
 128 func main() {
 129         seen := make(map[util.CompactPeer]struct{})
 130 getPeers:
 131         for {
 132                 ps, err := s.Announce(*infoHash, 0, false)
 133                 if err != nil {
 134                         log.Fatal(err)
 135                 }
 136         values:
 137                 for {
 138                         select {
 139                         case v, ok := <-ps.Peers:
 140                                 if !ok {
 141                                         break values
 142                                 }
 143                                 log.Printf("received %d peers from %x", len(v.Peers), v.NodeInfo.ID)
 144                                 for _, p := range v.Peers {
 145                                         if _, ok := seen[p]; ok {
 146                                                 continue
 147                                         }
 148                                         seen[p] = struct{}{}
 149                                         fmt.Println((&net.UDPAddr{
 150                                                 IP:   p.IP[:],
 151                                                 Port: int(p.Port),
 152                                         }).String())
 153                                 }
 154                         case <-quitting:
 155                                 break getPeers
 156                         }
 157                 }
 158                 if *once {
 159                         break
 160                 }
 161         }
 162         if err := saveTable(); err != nil {
 163                 log.Printf("error saving node table: %s", err)
 164         }
 165 }