#!/usr/bin/ruby

#USAGE NOTES
#bundle exec ruby .\wif2ansible.rb TAS000000535469.xlsx

require 'roo'
require 'yaml'
require 'resolv'
require 'uri'
require 'socket'

OSCACHE = {}
UNREACHABLE_HOSTS = []


class Hash
  def fuzzy_find_first(find)
    select { |key, value| key.to_s.match(/#{find}/i) }
  end
  def find_first(find)
    select { |key, value| key.to_s.match(/#{find}/i) }.values.flatten.first
  end
  def find_custom(find)
    select { |key, value| key.to_s.match(/#{find}/i) }.values.flatten
  end
end

wif_file = ARGV[0]
command_sheet_name = ARGV[1] ? ARGV[1] : "flow app"
portwass_switches = ARGV[2] ? ARGV[2] : "" #optionally add switches to portwas like -n

begin
  XLSX = Roo::Spreadsheet.open(wif_file, only_visible_sheets: true) 
rescue
  excel_fix_guide = %{
  1. Delete Sheet: Diagram
  2. Delete Sheet: Notes and Exceptions
  3. Save as xls
  4. Close Excel completely
  5. Open xls file
  6. Run document inspect, delete hidden rows
  7. Save the xls as xlsx}
  puts excel_fix_guide
end

#class Common
def find_sheet_name(name)
  result = []
  XLSX.sheets.each do |sheet_name| 
    if sheet_name.scan(/#{name.gsub(' ', '.*')}/i).any? 
      result << sheet_name 
    end
  end
  return result
end

def remove_html(string)
  if string.class == String
    string.split(/\<.*?\>/)
     .map(&:strip)
     .reject(&:empty?)
     .join(' ')
     .gsub(/\s,/,',').gsub('*', '').strip
   else
    string
   end
end

def is_windows?(servername)
  if OSCACHE[servername] == 'win'
    return true
  elsif OSCACHE[servername] == 'lin'
    return false
  #elsif OSCACHE[servername] == false
  #  puts "Timeout for #{servername}, skipping..."
  #  return false
  end
  
  attempts = 0
  begin
    attempts+=1
    s = TCPSocket.new servername, 3389
    OSCACHE[servername] = 'win' if s
    return true if s
  rescue Errno::ECONNREFUSED
    puts "#{servername}: Port 3389 not open, #{servername} is not a windows server"
    s = nil
    begin; s = TCPSocket.new servername, 22; rescue;OSCACHE[servername] = false; UNREACHABLE_HOSTS << servername;end;
    OSCACHE[servername] = 'lin' if s
    return false
  rescue IO::TimeoutError
    retry unless attempts > 2
    puts "#{servername}: IO Timeout to #{servername}. You may not be connected to the correct EDC. Please connect your VPN or run from a JUMP server in the correct EDC"
    OSCACHE[servername] = false
    return false
  end
end

def is_linux?(servername)
  if OSCACHE[servername] == 'lin'
    return true
  elsif OSCACHE[servername] == 'win'
    return false
  #elsif OSCACHE[servername] == false
  #  puts "Timeout for #{servername}, skipping..."
  #  return false
  end

  attempts = 0
  begin
    attempts+=1
    s = TCPSocket.new servername, 22
    OSCACHE[servername] = 'lin'
    return true if s
  rescue Errno::ECONNREFUSED
    puts "#{servername}: Port 22 not open, #{servername} is not a linux server"
    s = nil
    begin; s = TCPSocket.new servername, 3389; rescue;OSCACHE[servername] = false; UNREACHABLE_HOSTS << servername;end;
    OSCACHE[servername] = 'win' if s
    return false
  rescue IO::TimeoutError
    retry unless attempts > 2
    puts "#{servername}: IO Timeout to #{servername}. You may not be connected to the correct EDC. Please connect your VPN or run from a JUMP server in the correct EDC"
    OSCACHE[servername] = false
    return false
  end
end

def select_value_from_row(row, column)
  row.each{|k,v| return [k,cell_value_to_array(v)] if not v.nil? and not k.nil? and k.gsub("\n", '').scan(/#{column.gsub(' ', '.*')}/i).any? }
end

def cell_value_to_array(value)
  value.to_s.split(/[\n, " ", ","]/).compact.keep_if{|a| a.gsub(' ', '') != "" }
end
#end #Common

#class Flow

def is_empty_or_example_flow_row?(row)
  nil_count = 0
  nil_count_limit = 5
  row.each do |k,v|
    nil_count = nil_count + 1 if v.nil? 
  end
  if nil_count >= nil_count_limit or row.first[1].class == String or row.first[1].nil?
    return true
  else
    return false 
  end
end

def get_all_rows_and_find_headers(name)
  flow_header_items = [/flow/i, /source/i, /destination/i, /public/i, /ip/i, /private/i, /port/i]
  begin
    sheet_name = find_sheet_name(name).last
    puts "Using sheet: #{sheet_name}"
    XLSX.sheet(sheet_name).parse(header_search: flow_header_items)
  rescue Roo::HeaderRowNotFoundError
    sheet_name = find_sheet_name(name).first
    puts "ERROR: POSSIBLE EXTRA SHEET, trying to fix...\nUsing sheet: #{sheet_name}"
    begin
      XLSX.sheet(sheet_name).parse(header_search: flow_header_items)
    rescue Roo::HeaderRowNotFoundError
      puts "ERROR: Flow sheet table header names are incorrect in provided WIF. This script is looking for the following words: #{flow_header_items}. Fix this in Excel and use your modified WIF file."
      puts "Specifically, I want to see 'Source Public IP' 'Source Private IP' 'Destination Public IP' 'Destination Private IP'. Add columns if they have been deleted."
      exit 
    end
  end
end

def remove_udp_ports(value)
  value = value.to_s if value.class == Array
  value.to_s.gsub(/\d{2,5}.{1}udp/i, '')
end

def parse_ports(value)
  value = remove_udp_ports(value)
  value = value.to_s if value.class == Array
  #port_ranges = value.scan(/\d{2,5}-\d{2,5}|\d{2,5} - \d{2,5}/) 
  port_numbers = [value.scan(/\d{2,5}/)].flatten.map{|port| port.to_i}
  #if port_ranges.any?
  #  port_numbers = [(port_numbers + port_ranges)].flatten.compact.map{|range| range.to_s.split('-') }
  #  port_numbers = [port_numbers].flatten!.uniq!.map{|port| port.to_i}
  #end
  if value.scan(/any|all/i).any? && !port_numbers.any?
    return [22,3389,80,443,3306,5432,8443,60000] #return some frequently used ports if they requested all/any
  else
    return port_numbers
  end
end


def is_empty_or_example_flow_row?(row)
  all_source_ips = []
  all_source_ips << select_value_from_row(row, 'source private ip')[1]
  all_source_ips << select_value_from_row(row, 'source public ip')[1]
  all_source_ips = all_source_ips.flatten.compact

  all_destination_ips = []
  all_destination_ips << select_value_from_row(row, 'destination private ip')[1]
  all_destination_ips << select_value_from_row(row, 'destination public ip')[1]
  all_destination_ips = all_destination_ips.flatten.compact

  if !all_source_ips.any? or !all_destination_ips.any? or row.first[1].class == String or row.first[1].nil?
    return true
  else
    return false 
  end
end

#only flows that contain source ip and destination ip
def testable_flow_rows(sheet_name)
  begin 
    get_all_rows_and_find_headers(sheet_name).keep_if{|a| !is_empty_or_example_flow_row?(a) }
  rescue TypeError
    puts "ERROR: Problem accessing sheet with '#{sheet_name}' in the name. Does this sheet exist?"
    exit
  end
end

def cleanup_flow_formatting(rows)
  rows.map do |row|
    result = {}
    row.each do |k,v|
      cleaned_key = remove_html(k).to_s.gsub(" ", "_").gsub("\n", "_")
      if cleaned_key.scan(/port/i).any?
        cleaned_value = parse_ports(remove_html(v))
      elsif v.class == Integer
        cleaned_value = v
      elsif v.to_s.scan(/\b[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\b/).any? #ip addresses
        cleaned_value = cell_value_to_array(remove_html(v)).to_s.scan(/\b[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\b/)
      else
        cleaned_value = cell_value_to_array(remove_html(v))
      end
      result[cleaned_key] = cleaned_value
    end
    result
  end
end
#end #Flow


#class Server
def is_empty_or_example_server_row?(row)
  if select_value_from_row(row, 'cpu')[1].nil? || select_value_from_row(row, 'ram')[1].nil? || select_value_from_row(row, 'server reference')[1] == "Example"
    return true
  else
    return false 
  end
end

  def find_header_in_server_sheet
    sheet_name = find_sheet_name('server').first
    puts "Using sheet: #{sheet_name}"
    XLSX.sheet(sheet_name).parse(header_search: [/reference/i, /type/i, /alias/i, /platform/i, /middleware/i ]).first.map{|k,v| k}
  end

  def all_servers
    sheet_name = find_sheet_name('server').first
    XLSX.sheet(sheet_name).parse(header_search: [/reference/i, /type/i, /alias/i, /platform/i, /middleware/i ]).keep_if{|a| !is_empty_or_example_server_row?(a) }
  end

  def reference_to_ip(server_reference)
    matching = []
    all_servers.map do |row|
      this_reference = select_value_from_row(row, 'reference')
      if this_reference[1][0].scan(/#{server_reference}/i).any?
        matching << row
      end
    end
    return [] if !matching.any?
    return matching.first.keep_if{|k,v| k.scan(/ip.*address/i).any? if !k.nil? and !v.nil? }.map{|k,v| v}.sort!
  end


  def to_mgt_ip(name)
    begin
      fqdn = Resolv.getname(name)
      mgt_dns = fqdn.split('.').first + '.ds.gc.ca'
      mgt_ip = Resolv.getaddress(mgt_dns)
    rescue Resolv::ResolvError
      begin
        puts "#{name} not found in ds.gc.ca, checking pre-ds.gc.ca..."
        fqdn = Resolv.getname(name)
        mgt_dns = fqdn.split('.').first + '.pre-ds.gc.ca'
        mgt_ip = Resolv.getaddress(mgt_dns)
      rescue Resolv::ResolvError
        puts "#{name} is not a server OR no DNS entries exist in ds.gc.ca or pre-ds.gc.ca, skipping source..."
      end
    end
    #return mgt_ip
    return mgt_dns if mgt_ip.to_s.length > 4
  end


  def flows_by_host_to_ansible_inventory_yaml(flows_by_host)
    return {"all" => {"hosts" => flows_by_host}}.to_yaml
  end

#class Result
  def parse_portwass(stdout)
    stdout.scan(/^(\d{2,5}):\ (\w{4})/)
  end
#end

ansible_inventory_hash ={}

ansible_tasks = []
flows = []
if ARGV[1]
  begin
    flows = cleanup_flow_formatting(testable_flow_rows(ARGV[1]))
  rescue
    puts "INFO: unable to parse sheet containing '#{ARGV[1]}'"
    exit
  end
end
begin
  flows << cleanup_flow_formatting(testable_flow_rows('flow app'))
rescue
  puts "INFO: unable to parse sheet containing 'flow app'"
end
begin
  flows << cleanup_flow_formatting(testable_flow_rows('flow man'))
rescue
  puts "INFO: unable to parse sheet containing 'flow man'"
end
flows.flatten!

results = {}
mgt_ip_list = []
failed_portwass_cmds = []
flows_count = 0
flows.each do |flow|
  puts "\n\n#{'#'*8} Parsing flow number: #{flow['Flow_#']} #{'#'*8}"
  if results[flow['Flow_#']].nil?
    results[flow['Flow_#']] = {}
    results[flow['Flow_#']]['connections'] = [] 
  end
  if flow["Source_Public_IP"].nil? || flow["Source_Private_IP"].nil? || flow["Destination_Private_IP"].nil? || flow["Source_Public_IP"].nil?
    puts "ERROR IN SPREADSHEET:\n\n Please ensure there are columns named (case sensitive) \"Source Private IP\", \"Source Public IP\", \"Destination Private IP\", \"Destination Public IP\"\n\n Please update the names of columns and possibly add empty columns with these names if they have been combined."
    exit
  elsif flow["Flow_#"].nil?
    puts "ERROR IN SPREADSHEET:\n\n Please ensure the Flow # column is named (case sensitive) \"Flow #\""
    exit
  end
  flow_src_ips =  flow["Source_Public_IP"].any? ? flow["Source_Public_IP"] : flow["Source_Private_IP"]
  flow_src_ips.each do |src_ip|
    mgt_ip = to_mgt_ip(src_ip)
    mgt_ip_list << mgt_ip
    flow_dst_ips = flow["Destination_Public_IP"].any? ? flow["Destination_Public_IP"] : flow["Destination_Private_IP"]
    puts "Destination IPs empty for flow #{flow['Flow_#']}, skipping" && next if flow_dst_ips.nil? || !flow_dst_ips.any?
    flow_dst_ips.each do |dst_ip|

      

      if mgt_ip.to_s.length < 3
        puts "skipping #{mgt_ip} #{src_ip} as I don't think it's a windows/linux server" 
        next
      end
      if ansible_inventory_hash[mgt_ip].nil?
        if is_linux? mgt_ip
          ansible_inventory_hash.merge!({mgt_ip => {"flows" => []}  })
        elsif is_windows? mgt_ip
          ansible_inventory_hash.merge!({mgt_ip => {"ansible_connection" => "winrm", "ansible_winrm_transport" => "ntlm", "ansible_winrm_port" => 5985, "flows" => []}  })
        else
          #add to list fo unreachable hosts to output at end
          UNREACHABLE_HOSTS << mgt_ip
          UNREACHABLE_HOSTS << src_ip
        end

        #puts ansible_inventory_hash.to_yaml
      end
      begin
        a ={ "flow_number" => flow['Flow_#'], "dest" => dst_ip, "ports" => [flow.find_custom("Port")].flatten.uniq } 
        ansible_inventory_hash[mgt_ip]["flows"] << a
        puts "#{mgt_ip} : #{a}"
        flows_count +=1
      rescue NoMethodError
        puts "SKIPPING ERROR: #{mgt_ip} flow #{flow['Flow_#']}"
      end



    end#dst
  end#src

  #end

end#flows


puts ansible_inventory_hash
puts flows_by_host_to_ansible_inventory_yaml(ansible_inventory_hash)



ansible_inventory_filename = "#{File.basename(wif_file, "xlsx")[0..22]}_inventory_#{Time.now.strftime("%d-%m-%Y_%H.%M")}.yml"
File.open ansible_inventory_filename, 'w' do |file|
  file.write flows_by_host_to_ansible_inventory_yaml(ansible_inventory_hash)
end



puts "="*20
puts "Source servers found: #{ansible_inventory_hash.count}"

puts "Total connections: #{flows_count}"
puts "Generated inventory: #{ansible_inventory_filename}"
puts "="*20

if UNREACHABLE_HOSTS.compact.uniq.any?
  puts "The following [#{UNREACHABLE_HOSTS.compact.uniq.count}] servers could not be reached. Either they are in another datacentre or arent windows/linux servers:"
  puts UNREACHABLE_HOSTS.compact.uniq.join("\n")
end