DHTShellExample.java

import com.frostwire.jlibtorrent.AlertListener;
import com.frostwire.jlibtorrent.Dht;
import com.frostwire.jlibtorrent.DhtRoutingBucket;
import com.frostwire.jlibtorrent.Downloader;
import com.frostwire.jlibtorrent.Entry;
import com.frostwire.jlibtorrent.Pair;
import com.frostwire.jlibtorrent.Session;
import com.frostwire.jlibtorrent.Sha1Hash;
import com.frostwire.jlibtorrent.TcpEndpoint;
import com.frostwire.jlibtorrent.alerts.Alert;
import com.frostwire.jlibtorrent.alerts.AlertType;
import com.frostwire.jlibtorrent.alerts.DhtBootstrapAlert;
import com.frostwire.jlibtorrent.alerts.DhtImmutableItemAlert;
import com.frostwire.jlibtorrent.alerts.DhtMutableItemAlert;
import com.frostwire.jlibtorrent.alerts.DhtPutAlert;
import com.frostwire.jlibtorrent.alerts.DhtStatsAlert;
import com.frostwire.jlibtorrent.alerts.ListenFailedAlert;
import com.frostwire.jlibtorrent.alerts.ListenSucceededAlert;
import com.frostwire.jlibtorrent.alerts.PortmapErrorAlert;
import com.sun.org.apache.xerces.internal.impl.dv.util.HexBin;

import org.apache.commons.io.FileUtils;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Scanner;

/**
 * Created by onkar on 7/8/16.
 */
public class DHTShellExample {
  public static void main(String[] args) throws IOException {
    DataClass.loadLibrary("/lib/x86_64/libjlibtorrent.so");
    final Session s = new Session();

    AlertListener mainListener = new AlertListener() {
      public int[] types() {
        return null;
      }

      public void alert(Alert<?> alert) {
        AlertType type = alert.type();
        if (type == AlertType.LISTEN_SUCCEEDED) {
          ListenSucceededAlert a = (ListenSucceededAlert) alert;
          log(a.message());
        }

        if (type == AlertType.PORTMAP_ERROR) {
          PortmapErrorAlert a = (PortmapErrorAlert) alert;
          log(a.message());
        }

        if (type == AlertType.DHT_BOOTSTRAP) {
          DhtBootstrapAlert bootstrapAlert = (DhtBootstrapAlert) alert;
          log(bootstrapAlert.message());
        }
        if (type == AlertType.LISTEN_FAILED) {
          ListenFailedAlert a = (ListenFailedAlert) alert;
          log(a.message());
        }

        if (type == AlertType.DHT_PUT) {
          DhtPutAlert a = (DhtPutAlert) alert;
          /*System.out.println("PK = "+ HexBin.encode(a.publicKey()));
          System.out.println("Sig " + HexBin.encode(a.signature()));*/
          log(a.message());
        }

        if (type == AlertType.DHT_MUTABLE_ITEM) {
          DhtMutableItemAlert a = (DhtMutableItemAlert) alert;
          /*System.out.println(a.item().toString());
          System.out.println("PK = "+ HexBin.encode(a.key()));
          System.out.println("Sig " + HexBin.encode(a.signature()));*/
          Entry item = a.item();
          log(item.toString());
          log(a.message());
        }
        // ;
        if (type == AlertType.DHT_IMMUTABLE_ITEM) {
          DhtImmutableItemAlert a = (DhtImmutableItemAlert) alert;
          //System.out.println(a.getItem().toString());
          log(a.message());
        }

        if (type == AlertType.DHT_STATS) {
          DhtStatsAlert a = (DhtStatsAlert) alert;
          ArrayList<DhtRoutingBucket> dhtRoutingBuckets = a.routingTable();
          for (int i = 0; i < dhtRoutingBuckets.size(); i++) {
            System.out.println(dhtRoutingBuckets.get(i).numNodes());
          }
          long nodes = s.getStats().dhtNodes();// TODO: restore this
          log("DHT contains " + nodes + " nodes");
        }

      }
    };

    s.addListener(mainListener);

    Dht dht = new Dht(s);
    s.postDHTStats();
    Downloader downloader = new Downloader(s);

    try {
      File f = new File("dht_shell.dat");
      if (f.exists()) {
        byte[] data = FileUtils.readFileToByteArray(f);
        s.loadState(data);
      }
    } catch (Throwable e) {
      log(e.getMessage());
    }

    s.addDHTRouter(new Pair<String, Integer>("router.utorrent.com", 6881));
    s.addDHTRouter(new Pair<String, Integer>("router.bittorrent.com", 6881));
    s.addDHTRouter(new Pair<String, Integer>("dht.transmissionbt.com", 6881));
    s.addDHTRouter(new Pair<String, Integer>("router.bitcomet.com", 6881));
    s.addDHTRouter(new Pair<String, Integer>("dht.aelitis.com", 6881));

    Scanner in = new Scanner(System.in);
    while (true) {
      System.out.print("$ ");

      String line = in.nextLine().trim();

      if (is_quit(line)) {
        quit(s);
      } else if (is_put(line)) {
        put(dht, line);
      } else if (is_get(line)) {
        get(dht, line);
      } else if (is_get_peers(line)) {
        get_peers(dht, line);
      } else if (is_announce(line)) {
        announce(dht, line);
      } else if (is_mkeys(line)) {
        mkeys(line);
      } else if (is_mput(line)) {
        mput(dht, line);
      } else if (is_mget(line)) {
        mget(s, line);
      } else if (is_magnet(line)) {
        magnet(downloader, line);
      } else if (is_count_nodes(line)) {
        count_nodes(s);
      } else if (is_invalid(line)) {
        invalid(line);
      }
    }
  }

  private static void print(String s, boolean dollar) {
    System.out.println();
    if (dollar) {
      System.out.print("$log ");
    }
    System.out.println(s);
  }

  private static void print(String s) {
    print(s, false);
  }

  private static void log(String s) {
    print(s, true);
  }

  private static boolean is_quit(String s) {
    s = s.split(" ")[0];
    return s.equals("quit") || s.equals("exit") || s.equals("stop");
  }

  private static void quit(Session s) {
    print("Exiting...");
    byte[] data = s.saveState();
    try {
      FileUtils.writeByteArrayToFile(new File("dht_shell.dat"), data);
    } catch (Throwable e) {
      print(e.getMessage());
    }
    s.abort();
    System.exit(0);
  }

  private static boolean is_put(String s) {
    return s.startsWith("put ");
  }

  private static void put(Dht dht, String s) {
    String data = null;
    try {
      data = FileUtils.readFileToString(new File("note.json"));
    } catch (IOException e) {
      e.printStackTrace();
    }
    String sha1 = dht.put(new Entry(data)).toString();
    print("Wait for completion of put for key: " + sha1);
  }

  private static boolean is_get(String s) {
    return s.startsWith("get ");
  }

  private static void get(Dht dht, String s) {
    String sha1 = s.split(" ")[1];
    print("Waiting a max of 20 seconds to get data for key: " + sha1);
    Entry data = dht.get(new Sha1Hash(sha1), 20);
    print(data.toString());
  }

  private static boolean is_get_peers(String s) {
    return s.startsWith("get_peers ");
  }

  private static void get_peers(Dht dht, String s) {
    String sha1 = s.split(" ")[1];
    print("Waiting a max of 20 seconds to get peers for key: " + sha1);
    ArrayList<TcpEndpoint> peers = dht.getPeers(new Sha1Hash(sha1), 20);
    print(peers.toString());
  }

  private static boolean is_announce(String s) {
    return s.startsWith("announce ");
  }

  private static void announce(Dht dht, String s) {
    String sha1 = s.split(" ")[1];
    dht.announce(new Sha1Hash(sha1), 9000, 0);
    print("Wait for completion of announce for key: " + sha1);
  }

  private static boolean is_mkeys(String s) {
    return s.startsWith("mkeys");
  }

  private static void mkeys(String s) {
    byte[][] keys = Dht.createKeypair();
    String msg = "Save this key pair\n";
    msg += "Public: " + HexBin.encode(keys[0]) + "\n";
    msg += "Private: " + HexBin.encode(keys[1]) + "\n";
    print(msg);
  }

  private static boolean is_mput(String s) {
    return s.startsWith("mput ");
  }

  private static void mput(Dht dht, String s) {
    String[] arr = s.split(" ");
    byte[] publicKey = HexBin.decode(arr[1]);
    byte[] privateKey = HexBin.decode(arr[2]);
    String data = null;
    try {
      data = FileUtils.readFileToString(new File("note.json"));
    } catch (IOException e) {
      e.printStackTrace();
    }
    byte[] salt = arr[4].getBytes();
    dht.mput(publicKey, privateKey, new Entry(data), salt);
    print("Wait for completion of mput for public key: " + arr[1] + " ");
    print("Data is " + data);
  }

  private static boolean is_mget(String s) {
    return s.startsWith("mget ");
  }

  private static void mget(Session session, String s) {
    String[] arr = s.split(" ");
    byte[] publicKey = HexBin.decode(arr[1]);
    byte[] salt = arr[2].getBytes();
    print("Waiting a max of 20 seconds to get mutable data for public key: " + arr[1]);
    /*Dht.MutableItem data = dht.mget(publicKey, salt, 50);
    if (data != null)
      print(data.item.toString());*/
    session.dhtGetItem(publicKey, salt);
  }

  private static boolean is_magnet(String s) {
    return s.startsWith("magnet ");
  }

  private static void magnet(Downloader downloader, String s) {
    String sha1 = s.split(" ")[1];
    String uri = "magnet:?xt=urn:btih:" + sha1;
    print("Waiting a max of 20 seconds to fetch magnet for sha1: " + sha1);
    byte[] data = downloader.fetchMagnet(uri, 20);
    print(Entry.bdecode(data).toString());
  }

  private static boolean is_count_nodes(String s) {
    return s.startsWith("count_nodes");
  }

  private static void count_nodes(Session s) {
    s.postDHTStats();
  }

  private static boolean is_invalid(String s) {
    return !s.isEmpty();
  }

  private static void invalid(String s) {
    print("Invalid command: " + s + "\n" + "Try ? for help");
  }


}