# Copyright (c) 2022 Marek Kuethe
# License: WTFPL

# This script pings all Yggdrasil nodes and then displays a recommendation of
# nodes. The script is assumed to be in the public-peers folder. If not, a
# different folder can be given as an argument. Nodes that cannot be reached via
# TLS or TCP are omitted. Each node is pinged 20 times over TCP. There is a 0.2
# second pause after each ping.

# Require gem net-ping

require "net/ping"

$pinger = Net::Ping::TCP

def yggdrasil_node_list dir = "./"

  list = {}
  max_len_region = 0
  max_len_node = 0

  Dir["#{dir}*/*.md"].each { |file|
    if File.file?(file)
      cnt = File.read file

      path = file.split("/")
      item = "#{path[-2]}/#{path[-1].split(".")[0]}"
      max_len_region = item.length if item.length > max_len_region
      list[item] = cnt.scan(/\`(tls|tcp|socks):\/\/([a-z0-9\.\-\:\[\]]+):([0-9]+)\`/).to_a
      list[item].each { |node|
        len_node = "#{node[1]}:#{node[2]}".length
        max_len_node = len_node if len_node > max_len_node
      }
      
      # regex by https://github.com/zhoreeq/peer_checker.py/blob/master/md_to_json.py
    end
  }
  
  return [list, max_len_node, max_len_region]

end

def ping_node node_info, count = 20, timeout = 5, pause = 0.2
  sum = 0.0
  fail = 0
  ping = $pinger.new node_info[1], node_info[2], timeout
  count.times {
    res = ping.ping
    if res
      sum += res
    else
      fail += 1
    end
    sleep pause
  }
  sum /= count
  
  return [sum, fail]
end

def collect_ping_times list, count = 20
  ping_res = {}
  threads = []
  # ping over socks is not supported
  protocols = ["tcp", "tls"]
  
  puts "Collecting ping times..."
  list[0].each_pair { |region, nodes|
    # open a new thread for every country / subregion
    threads << Thread.new(region, nodes) { |region, nodes|
      puts "Pinging nodes from region #{region}"
      nodes.each { |node_info|
        if protocols.include? node_info[0]
          puts "Ping node #{node_info[1]} (#{node_info[0]})"
          res = ping_node node_info, count
          puts "count=#{count}  avg=#{res[0].round 3}  fail=#{res[1]}"
          res << region
          ping_res[node_info] = res
        end
      }
    }
  }
  
  threads.each.with_index { |th, index|
    puts "Waiting for thread #{index + 1}..."
    th.join
  }
  
  return ping_res
end

def display_ping_times list, ping_res
  node_with_fail = {}
  node_without_fail = {}
  
  name_len = list[1] + 2
  region_len = list[2] + 2
  
  # split hash into nodes with and without ping fails
  ping_res.each_pair { |node, res|
    if res[1] == 0
      node_without_fail[node] = res
    else
      node_with_fail[node] = res
    end
  }
  
  # sort by ping time
  node_with_fail = node_with_fail.sort_by { |node, res|
    res[0]
  }
  node_without_fail = node_without_fail.sort_by { |node, res|
    res[0]
  }
  
  puts "=" * 43
  puts
  puts "Nodes with ping fails:"
  node_with_fail.each { |node_res|
    node_name = "#{node_res[0][1]}:#{node_res[0][2]}"
    name_padding = " " * (name_len - node_name.length)
    region = "region=#{node_res[1][2]}"
    region_padding = " " * (region_len - node_res[1][2].length)
    data = "fail=#{node_res[1][1]}\tavg=#{node_res[1][0].round 3}"
    puts "#{node_name}#{name_padding} (#{node_res[0][0]})   #{region}#{region_padding}#{data}"
  }
  
  puts
  puts "-" * 43
  puts
  puts "Nodes without ping fails:"
  node_without_fail.each { |node_res|
    node_name = "#{node_res[0][1]}:#{node_res[0][2]}"
    name_padding = " " * (name_len - node_name.length)
    region = "region=#{node_res[1][2]}"
    region_padding = " " * (region_len - node_res[1][2].length)
    data = "avg=#{node_res[1][0].round 3}"
    puts "#{node_name}#{name_padding} (#{node_res[0][0]})   #{region}#{region_padding}#{data}"
  }
  
  puts
  puts "-" * 43
  puts
  puts "I would recommend you to add the following three nodes:"
  puts "  Peers: ["
  node_without_fail.first(3).each { |node, res|
    puts "    #{node[0]}://#{node[1]}:#{node[2]}  \# #{res[2]}"
  }
  puts "  ]"
  puts
  
  puts "If you want more choice, here are the twenty nodes with the lowest ping time:"
    node_without_fail.first(20).each { |node, res|
    puts "    #{node[0]}://#{node[1]}:#{node[2]}  \# #{res[2]}"
  }
end

def test_all dir = "./"
  list = yggdrasil_node_list dir
  display_ping_times list, collect_ping_times(list)
end

dir = "./"

if ARGV[0] && File.directory?(ARGV[0])
  dir = ARGV[0]
end

test_all dir