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