#!/usr/bin/ruby
require 'fileutils'
require 'faker'
require 'logger'
require 'net/http'
require 'optparse'
require './labstats'

raise RuntimeError, 'This script has to be run with root permissions' unless Process.uid == 0

@conf = {
  :do_configure_peering => true,
  :do_create_database => true,
  :do_create_pgp_keys => true,
  :do_initial_setup => true,
  :do_install_hockeypuck => true,
  :do_install_attest_ck => true,
  :do_pgp_poison => true,
  :do_sign_pgp_keys => true,
  :guest_name => 'hkp%02d',
  :guest_dir => '/var/lib/lxc/%s',
  :hkp_pkg_base => 'hockeypuck-pcic_1.0-1_amd64.deb',
  :hkp_pkg_modif => 'hockeypuck-pcic_1.0-2_amd64.deb',
  :install_pkgs => %w(eatmydata screen ssmtp postgresql sudo git bzr mercurial
                      golang sq openssh-server libcapture-tiny-perl),
  :install_pkgs_attest => %w(eatmydata screen ssmtp netcat-openbsd git cargo
                             pkg-config libssl3 libssl-dev clang llvm
                             nettle-dev openssh-server libcapture-tiny-perl),
  :ip_gateway => '10.0.3.1',
  :ip_network => '10.0.3.%d/24',
  :log_prio => :debug,
  :log_to => 'build_img.log',
  #:log_to => STDOUT,
  :modified_hkp => true,
  :num_hosts => 5,
  :peer_prob => 0.3,
  # How likely is it that a signature will be attested? (in
  # relation with the total signatures created at
  # pgp_sign_prob)
  :pgp_attested_sign_prob => 0.3,
  # How likely is that a given key will sign any other key
  :pgp_sign_prob => 0.01,
  :pgp_keydir_filename => 'pgp_keydir.txt',
  :pgp_keys_dir => 'generated_pgp_keys',
  # Total keys generated will be pgp_keys_per_host × num_hosts
  :pgp_keys_per_host => 100,
  :pgp_keys_to_poison => 5,
  :pgp_attacker_keys => 1000,
  :use_eatmydata => true
}

OptionParser.new do |opts|
  opts.banner = 'Usage: build_img.rb [options]'
  opts.on('-h', '--help', 'Prints this help'){puts opts;exit 0}
  opts.on('-n NUM', '--num_hosts', 'Number of HKP servers to build') { |n|
    @conf[:num_hosts] = n.to_i
  }
  opts.on('-A', '--no-attest-ck',
          'Skip attestation checker setup') {
    @conf[:do_initial_attest_ck] = false
  }
  opts.on('-C', '--no-config-peering',
          'Skip the creation of a peering network') {
    @conf[:do_configure_peering] = false
  }
  opts.on('-D', '--no-create-db',
          'Skip PostgreSQL DB initialization') {
    @conf[:do_create_database] = false
  }
  opts.on('-H', '--no-install-hkp',
          'Skip downloading and installing Hockeypuck') {
    @conf[:do_install_hockeypuck] = false
  }
  opts.on('-I', '--no-initial',
          'Skip initial container creation') {
    @conf[:do_initial_setup] = false
  }
  opts.on('-M', '--no-modified-prot',
          'Install Hockeypuck without the protocol modification') {
    @conf[:modified_hkp] = false
  }
  opts.on('-P', '--no-create-pgp',
          'Skip the creation of OpenPGP keys') {
    @conf[:do_create_pgp_keys] = false
  }
  opts.on('-S', '--no-sign-keys',
          'Skip the signature of OpenPGP keys') {
    @conf[:do_sign_pgp_keys] = false
  }
  opts.on('-X', '--no-poison-pgp',
          'Skip the poisoning of OpenPGP keys') {
    @conf[:do_pgp_poison] = false
  }
end.parse!


def build_guest(guest_num)
  initial_setup(guest_num) if @conf[:do_initial_setup]
  create_database(guest_num) if @conf[:do_create_database]
  install_hockeypuck(guest_num) if @conf[:do_install_hockeypuck]
  configure_peering(guest_num) if @conf[:do_configure_peering]
end

def initial_setup(guest_num)
  guest = guest_name_for(guest_num)
  # Create the LXC base install
  @log.info("Installing base system packages for #{guest}")
  cmd = %W(lxc-create -t debian -n #{guest} -- -r bullseye
           --packages=#{@conf[:install_pkgs].join(',')} )
  cmd.unshift('eatmydata') if @conf[:use_eatmydata]
  system(*cmd)

  # For reliability and predictable network addressing: Set up the
  # network configuration in the lxc configuration, comment it from
  # the system
  @log.info("Configuring network for #{guest}")
  File.open( File.join(@conf[:guest_dir] % guest, 'config'),
             'a' ) do |conf|
    conf.puts "lxc.net.0.ipv4.address = #{ip_and_netmask_for(guest_num)}"
    conf.puts "lxc.net.0.ipv4.gateway = #{@conf[:ip_gateway]}"
  end
  net_conf = File.open(filename_in_guest(guest, '/etc/network/interfaces'), 'r'
                      ).readlines.map {|lin| lin="# #{lin}" if lin =~ /eth0/}
  File.open(filename_in_guest(guest, '/etc/network/interfaces'), 'w') do |net|
    net.puts(net_conf)
  end

  # Copy the "glue script" to interface between Hockeypuck and Sequoia
  File.copy_stream('filter_certs',
                   filename_in_guest(guest, '/usr/local/bin/filter_certs')
                  )
  File.chmod(0755, filename_in_guest(guest, '/usr/local/bin/filter_certs'))

  @log.info("Starting up container #{guest}")
  # Wait for things to settle (is eatmydata a good idea or not?) and
  # start the guest system up
  sleep(2)
  system('lxc-start', '-n', guest)
end

def attest_ck_setup()
  guest = 'attest';
  @log.info('Building attestation checker')
  cmd = %W(lxc-create -t debian -n #{guest} -- -r bookworm
           --packages=#{@conf[:install_pkgs_attest].join(',')})
  cmd.unshift('eatmydata') if @conf[:use_eatmydata]
  system(*cmd)

  # 'sq' is compiled from the modified version of Sequoia's repository
  File.copy_stream('sq', filename_in_guest(guest, '/usr/local/bin/sq'))
  File.chmod(0755, filename_in_guest(guest, '/usr/local/bin/sq'))

  @log.info("Configuring attestation checker's network")
  File.open( File.join(@conf[:guest_dir] % guest, 'config'),
             'a' ) do |conf|
    # 240 → 10.0.3.250
    conf.puts "lxc.net.0.ipv4.address = #{ip_and_netmask_for(240)}"
    conf.puts "lxc.net.0.ipv4.gateway = #{@conf[:ip_gateway]}"
  end
  net_conf = File.open(filename_in_guest(guest, '/etc/network/interfaces'), 'r'
                      ).readlines.map {|lin| lin="# #{lin}" if lin =~ /eth0/}
  File.open(filename_in_guest(guest, '/etc/network/interfaces'), 'w') do |net|
    net.puts(net_conf)
  end

  # "Glue script" to check with Sequoia the keys sent by the Sequoia servers
  File.copy_stream('tcp_to_sq.pl',
                   filename_in_guest(guest, '/usr/local/bin/tcp_to_sq'))
  File.chmod(0755, filename_in_guest(guest, '/usr/local/bin/tcp_to_sq'))
  File.copy_stream('tcp_to_sq.service',
                   filename_in_guest(guest, '/etc/systemd/system/tcp_to_sq.service'))

  @log.info("Starting up attestation container")
  # Start the guest system up
  system('lxc-start', '-n', guest)

  adduser = guest_running(guest,
                          %w(adduser --disabled-password --disabled-login
                             --gecos SequoiaHelper sq))
  Process.wait(adduser.pid)
  guest_running(guest, %w(systemctl enable tcp_to_sq))
  guest_running(guest, %w(systemctl start tcp_to_sq))
end

def create_database(guest_num)
  guest = guest_name_for(guest_num)
  @log.info("Creating database in #{guest}")
  # Wait until PostgreSQL is running, and create a HKP user and
  # database. Have to lxc-attach as "real" files in /run are not
  # exposed through the cgroups
  wait_for_psql = guest_running(guest,
      ['perl', '-e', 'until (-e "/var/run/postgresql/.s.PGSQL.5432") {print "."; sleep 1}'])
  Process.wait(wait_for_psql.pid)
  wait_for_psql.close

  psql = guest_running(guest,
                       %w(sudo -u postgres psql template1))
  psql.puts %q(CREATE USER hkp PASSWORD 'hkp-pcic-test';)
  psql.puts %q(CREATE DATABASE hkp OWNER hkp;)
  psql.close_write
  @log.debug '--- PostgreSQL results: ' + psql.read
  psql.close
end

def install_hockeypuck(guest_num)
  guest = guest_name_for(guest_num)
  @log.info("Installing Hockeypuck in #{guest}")

  # Copy and install the corresponding Hockeypuck package into the
  # container
  hkp_pkg = @conf[:modified_hkp] ?
              @conf[:hkp_pkg_modif] :
              @conf[:hkp_pkg_base]
  FileUtils.cp(hkp_pkg, filename_in_guest(guest, 'root'))
  dpkg = guest_running(guest,
                       [%W(dpkg -i /root/#{hkp_pkg})])
  Process.wait(dpkg.pid)
  @log.debug '--- dpkg results: ' + dpkg.read
  dpkg.close

  # Fix the configuration to use the database we just created
  fix_conf = guest_running(guest,
                           ['perl', '-p', '-i', '-e',
                            ('s/^# dsn/dsn/; s/h0ck3y/hkp-pcic-test/; ' +
                             's/INFO/DEBUG/'),
                            '/etc/hockeypuck-pcic/hockeypuck.conf'])
  Process.wait(fix_conf.pid)
  fix_conf.close

  # If our modifications are in place, the binaries have to be
  # rebuilt.
  if File.exists?(filename_in_guest(guest, '/usr/sbin/hockeypuck-rebuild'))
    rebuild = guest_running(guest, '/usr/sbin/hockeypuck-rebuild')
    Process.wait(rebuild.pid)
    rebuild.close
  end
end

def configure_peering(guest_num)
  guest = guest_name_for(guest_num)
  @log.info("Configuring peering between servers for #{guest}")
  conf = File.open(filename_in_guest(guest,
                   '/etc/hockeypuck-pcic/hockeypuck.conf'), 'a')

  @peers[guest_num].each_with_index do |peer, idx|
    # No vamos a hablar con nosotros mismos
    next if guest_num == idx
    # Si en el volado salió que con este no, pos no
    next unless peer

    peer_addr = ip_addr_for(idx)
    conf.puts ""
    conf.puts '[hockeypuck.conflux.recon.partner.%s]' %
              peer_addr.gsub(/\./,'-')
    conf.puts "httpAddr=\"#{peer_addr}:11371\""
    conf.puts "reconAddr=\"#{peer_addr}:11370\""
  end

  # Listos para iniciar el servicio de Hockeypuck! Monitoreamos hasta
  # que inicia correctamente (¡qué feo!)
  @log.info('About to start hockeypuck-pcic at %s' % guest)
  success = false
  times = 0
  while ! success
    sleep(1)
    times += 1
    start_hkp = guest_running(guest, ['systemctl', 'restart', 'hockeypuck-pcic'])
    pid, status = Process.wait2(start_hkp.pid)
    start_hkp.close

    success = true if status.exitstatus == 0
    @log.info('%s startup @%d: %d → %d (%s)' % [guest, times, pid, status,
                                       (success ? '✓' : '✗')])
  end
end

def create_pgp_keys(guest_num)
  @log.debug '--- Generating %d OpenPGP keys for guest %d' %
             [@conf[:pgp_keys_per_host], guest_num]

  @conf[:pgp_keys_per_host].times do |i|
    person = Faker::Name.name()
    mail = Faker::Internet.email(name: person)
    userid = '%s <%s>' % [person, mail]
    c_time = Faker::Time.backward(days: 1095).strftime('%Y-%m-%dT%H:%M:%S')
    exp_time = Faker::Time.forward(days: 1095).strftime('%Y-%m-%dT%H:%M:%S')

    Dir.mkdir(@conf[:pgp_keys_dir]) unless Dir.exists?(@conf[:pgp_keys_dir])
    keyfile = File.join(@conf[:pgp_keys_dir],
                        '%02d_%04d.%s.pgp' % [guest_num, i,
                                              person.split(/\s/).join('')]
                       ).gsub(/\.+/, '.')
    certfile = keyfile + '.cert'

    # Generación de la llave usando Sequoia (desde el anfitrión)
    system('sq', 'key', 'generate',
           '--userid', userid,
           '--creation-time', c_time,
           '--expires', exp_time,
           '--export', keyfile)
    system('sq', 'key', 'extract-cert', keyfile, '-o', certfile)

    # Registramos la llave y archivo en el archivo de
    # control/bitácora, para su posterior consumo a la hora de firmar
    # (omitiendo el nombre del directorio
    File.write(@conf[:pgp_keydir_filename],
               "#{File.basename(keyfile)}::#{userid}\n",
               mode: 'a+')

    # Enviamos la llave por HKP al servidor
    res = hkp_send_key(ip_addr_for(guest_num), File.read(certfile))
    print ' [%d-%d %s]' % [guest_num, i,
                           (res.nil? ? '✓' : '✗ ')]
  end

  # Damos una "patada" a Hockeypuck para que inicie la sincronización
  Process.wait(guest_running(guest_name_for(guest_num),
                             ['systemctl', 'restart', 'hockeypuck-pcic']
                            ).pid)
end

def gen_pgp_signatures()
  keyfiles = {}
  tot_sign = {:signatures => 0, :attestations => 0}
  File.read(@conf[:pgp_keydir_filename]).split(/\n/).each do |lin|
    lin =~ /^(.+)::(.+)$/ or
      @log.warn "Unexpected line parsing #{@conf[:pgp_keydir_filename]}: #{lin}"
    keyfiles[$1] = $2
  end

  keyfiles.keys.each do |key1|
    keyfiles.keys.each do |key2|
      # Sólo firmar con probabilidad ≥ pgp_sign_prob
      next unless rand <= @conf[:pgp_sign_prob]
      if tot_sign[:signatures] % 50 == 0
        print "\n%04d " % tot_sign[:signatures]
      end
      tot_sign[:signatures] += 1
      print '👍'
      signed = '%s-signs-%s.cert' % [key1, key2]

      uid = keyfiles[key2]
      @log.debug 'New signature %s → %s (%s)' % [key1, key2, uid]
      system('sq', 'certify',
             pgp_file_for(key1),
             pgp_file_for('%s.cert' % key2),
             uid,
             '--output', pgp_file_for(signed))

      # Atestiguar la firma si la probabilidad así lo marca
      if rand <= @conf[:pgp_attested_sign_prob]
        @log.debug '%s attests %s→%s signature' % [key2, key1, key2]
        print '✅'
        tot_sign[:attestations] += 1
        # Atestiguar y guardar en el mismo archivo (signed)
        signed_full_key = '%s.pgp' % signed
        system('sq', 'keyring', 'merge',
               pgp_file_for(key2),
               pgp_file_for(signed),
               '--output', pgp_file_for(signed_full_key))
        system('sq', 'key', 'attest-certifications',
               pgp_file_for(signed_full_key),
               '-o',
               pgp_file_for('%s.att' % signed_full_key))
        # Reescribimos pgp_file_for(signed), hay que especificar --force
        system('sq', '--force', 'key', 'extract-cert',
               '--output', pgp_file_for(signed),
               pgp_file_for('%s.att' % signed_full_key))
      end

      to_server = ip_addr_for( (0..@conf[:num_hosts]-1).to_a.sample )
      @log.debug 'Sending %s→%s key to server %s' % [key1, key2, to_server]
      hkp_send_key(to_server, File.read(pgp_file_for(signed)))
    end
  end
  return tot_sign
end

def pgp_poisoning(num)
  # Creamos 1000 llaves para montar el ataque. Por lo visto, Hockeypuck
  # ya se protege contra  pgp-poisoner.
  attack_dir = 'attack_keys'
  attack_pids = []
  Dir.mkdir(attack_dir) unless Dir.exists?(attack_dir)
  @log.debug('Creating %d poisoning keys' % @conf[:pgp_attacker_keys])
  print "\n⚗️"
  @conf[:pgp_attacker_keys].times do |i|
    print "🧪" if i % 50 == 0
    system('sq', '--force', 'key', 'generate',
           '--userid', 'Poisoning attacker key #%d' % i,
           '--export', File.join(attack_dir, i.to_s))
  end

  # Elegimos las víctimas aleatorias que nos indicaron para atacar
  File.read(@conf[:pgp_keydir_filename]).split(/\n/).sample(num).
    each_with_index do |v, idx|
    pid = fork()
    if ! pid.nil?
      # Proceso padre: Toma nota del hijo y lanza el siguiente ataque
      attack_pids << pid
      next
    end

    (victim, uid) = v.split('::')
    victim_dir = File.join(attack_dir, victim)
    Dir.mkdir(victim_dir) unless Dir.exists?(victim_dir)
    print " #{idx}☣️..."
    @log.debug('Attack %d: Poisoning key «%s»' % [idx, victim])
    # Generamos cada firma como un archivo independiente porque hkp
    # limita por tamaño máximo de POST
    @conf[:pgp_attacker_keys].times do |i|
      print '💀️'  if i%50 == 0
      to_server = ip_addr_for( (0..@conf[:num_hosts]-1).to_a.sample )
      IO.popen(['sq', 'certify', File.join(attack_dir, i.to_s),
                pgp_file_for(victim + '.cert'), uid]) do |sign|
        signed = sign.read
        hkp_send_key(to_server, signed)
      end
    end

    # El ataque se realiza desde un subproceso. Finalizamos el proceso
    # al terminar el ataque.
    exit(0)
  end

  @log.debug('%d attacks launched.' % num)
  attack_pids.each { |pid| Process.waitpid(pid) }
  @log.info('%d attacks finished.' % num)
end

def pgp_file_for(keyfile)
  return File.join(@conf[:pgp_keys_dir], keyfile)
end

def hkp_send_key(server, key_data)
  begin
    post_url = 'http://%s:11371/pks/add' % server
    res = Net::HTTP.post_form(URI(post_url), 'keytext' => key_data)
    return nil if res.code.to_i == 200
  rescue Errno::ECONNREFUSED
    @log.warn('Failed to send key to server %s: Connection refused' % server)
  rescue
    @log.warn('Error sending key to server %s: %s (%s)' %
              [server, res.code.to_i, res.message])
  end
end

def make_peer_list()
  peers=[]
  num = @conf[:num_hosts]

  # Crear un arreglo anidado vacío (más "bonito" trabajar con arreglos
  # cuadrados)
  num.times do |i|
    peers[i]=[]
    num.times do |j|
      peers[i][j] = nil
    end
  end

  # Llenar con la probabilidad especificada
  num.times do |i|
    num.times do |j|
      next if i==j
      r=rand()
      if r > (1 - @conf[:peer_prob])
        peers[i][j] = 1
        peers[j][i] = 1
      end
    end
  end
  return peers
end

def peering_matrix()
  num=0
  peers = []
  peers << 'Peering matrix for the test network:'
  peers << '==================================='
  peers << ''
  peers << "  " + @conf[:num_hosts].times.map{|i| i.to_s}.join('')
  peers << @peers.map do |lin|
    h = "#{num} "
    num+=1
    h + lin.map{|host| host ? 'X' : ' '}.join('')
  end

  return peers.join("\n")
end

def guest_name_for(guest_num)
  return @conf[:guest_name] % guest_num
end

def filename_in_guest(guest, target)
  return File.join(@conf[:guest_dir] % guest, 'rootfs', target)
end

def ip_and_netmask_for(guest_num)
  return @conf[:ip_network] % (guest_num + 10)
end

def ip_addr_for(guest_num)
  return ip_and_netmask_for(guest_num).gsub(/\/\d+/, '')
end

def guest_running(guest, cmd)
  lxc_cmd = %W(lxc-attach -n #{guest} --)
  lxc_cmd << cmd
  lxc_cmd.flatten!
  return IO.popen(lxc_cmd, 'r+')
end

@log = Logger.new(@conf[:log_to])
@log.level = @conf[:log_prio]
@log.info('--- Starting lab instantiation')
started_at = Time.now

# Creamos la lista de cómo se van a comunicar entre sí los servidores,
# para implementarla durante la instalación en cada archivo de
# configuración.
@peers = make_peer_list()

# Vaciamos el archivo de directorio de llaves. Lo manejamos como
# archivo externo para comunicarlo entre procesos (y sí, fuchi, sin
# mutex para protegerlo... Crucemos los deditos)
File.write(@conf[:pgp_keydir_filename], '')

# Lanzamos un proceso independiente para construir a cada uno de los
# servidores de la red, y uno adicional para el verificador de
# testificación
builders = []
(@conf[:num_hosts]+1).times do |guest_num|
  @log.info("--> #{guest_num} / #{@conf[:num_hosts]}")
  pid = Process.fork
  if pid
    builders << pid
  else
    if guest_num < @conf[:num_hosts]
      build_guest(guest_num)
      @log.info '*** Guest number %d finished building (%4.2f seconds)' %
                [guest_num, (Time.now - started_at)]
    else
      attest_ck_setup if @conf[:do_install_attest_ck]
      @log.info '*** Attestation checker finished building (%4.2f seconds)' %
                [guest_num, (Time.now - started_at)]
    end
    exit(0)
  end
end

# Esperamos a que terminen de construirse todos los servidores
builders.each {|pid| Process.wait(pid)}
@log.info '*** Laboratory built! (%4.2f seconds)' % (Time.now - started_at)
@log.info '    %d containers created (%s -- %s)' %
          [ @conf[:num_hosts],
            guest_name_for(0),
            guest_name_for(@conf[:num_hosts]-1)]

# Ya que el servidor de BD está corriendo, iniciamos los monitores de
# crecimiento de número de llaves y el histograma total de llaves
NumKeysMonitor.new(@conf[:guest_name], @conf[:num_hosts])
KeySizeHistogram.new(@conf[:guest_name], @conf[:num_hosts])

# Registramos la matriz de conexiones en cada uno de los servidores
# (para referencia más fácil al depurar)
if @conf[:do_configure_peering]
  peers = peering_matrix()
  @conf[:num_hosts].times do |i|
    peer_file = filename_in_guest( guest_name_for(i), '/root/peers.txt' )
    File.open(peer_file, 'w') {|f| f.puts peers}
  end
  @log.debug "Peering information:\n#{peers}"
end

# Allow for servers to start up before sending requests
sleep(5)
if @conf[:do_create_pgp_keys]
  pgp_creators = []
  @conf[:num_hosts].times do |guest_num|
    pid = Process.fork
    if pid
      pgp_creators << pid
    else
      create_pgp_keys(guest_num)
      exit(0)
    end
  end
  pgp_creators.each{|pid| Process.wait(pid)}
end

# Creamos las firmas de forma aleatoria entre las llaves generadas,
# siguiendo los parámetros :pgp_* de @conf
if @conf[:do_sign_pgp_keys]
  if ! @conf[:do_create_pgp_keys]
    @log.warn('Not creating key signatures: Keys not created, cannot' +
              'assume they exist.')
  end
  tot_sign = gen_pgp_signatures()
end

if @conf[:do_pgp_poison]
  @log.info('Poisoning and uploading %d generated keys' %
            @conf[:pgp_keys_to_poison])
  pgp_poisoning(@conf[:pgp_keys_to_poison])

  @log.info('All is done. Keep the program running to log synchronization.')
  print "\n\nTime elapsed since end of attack: "
  Thread.new do
    i=0
    print '    '
    while true
      i+=1
      print "\b\b\b\b\b%05d" % i
      sleep 1
    end
  end

  finished = false
  begin
    sleep(36000) # 10hr
    @log.info('Finished!.')
    finished = true
  ensure
    @log.info('Long sleep interrupted') unless finished
  end
end
exit 0