#!/usr/bin/env ruby
# frozen_string_literal: true

#
# gen-main.rb: Generate `main.c` for cavp-tests.  This script does the
# following:
#
# 1. Fetches the SHA-3 and SHAKE byte test vectors zip archives and
#    caches them in `./zip-cache`.
# 3. Verifies the SHA-256 digest of the downloaded archives.
# 3. Parses the response files (`.rsp`) in the downloaded zip archives
#    and extracts the SHA-3 and SHAKE test vectors.
# 4. Uses the test vectors to generate `main.c`.
# 5. Prints `main.c` to standard output.
#
# Usage:
#   cd tests/cavp-tests/
#
#   # download/cache archives, then generate main.c from them
#   ./gen-main.rb > main.c
#
#   # build and run tests
#   make && ./cavp-tests
#

# load libraries
require 'open-uri'
require 'openssl'
require 'uri'
require 'zip'

# number of elements per line of test data
PER_LINE = 128

# test vector sets
SETS = [{
  # URL to zip file containing test vectors for this set of algorithms
  url: 'https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Algorithm-Validation-Program/documents/sha3/sha-3bytetestvectors.zip',

  # expected sha256 digest of zip file
  hash: 'cd07701af2e47f5cc889d642528b4bf11f8b6eb55797c7307a96828ed8d8fc8c',

  # algorithms
  algos: [{
    name: 'SHA3-224', # algo name
    fn: 'sha3_224', # function prefix

    type: :hash, # function type
    size: 28, # output size, in bytes

    # response files
    rsp_files: %w{
      SHA3_224ShortMsg.rsp
      SHA3_224LongMsg.rsp
    },
  }, {
    name: 'SHA3-256', # algo name
    fn: 'sha3_256', # function prefix

    type: :hash, # function type
    size: 32, # output size, in bytes

    # response files
    rsp_files: %w{
      SHA3_256ShortMsg.rsp
      SHA3_256LongMsg.rsp
    },
  }, {
    name: 'SHA3-384', # algo name
    fn: 'sha3_384', # function prefix

    type: :hash, # function type
    size: 48, # output size, in bytes

    # response files
    rsp_files: %w{
      SHA3_384ShortMsg.rsp
      SHA3_384LongMsg.rsp
    },
  }, {
    name: 'SHA3-512', # algo name
    fn: 'sha3_512', # function prefix

    type: :hash, # function type
    size: 64, # output size, in bytes

    # response files
    rsp_files: %w{
      SHA3_512ShortMsg.rsp
      SHA3_512LongMsg.rsp
    },
  }],
}, {
  # URL to zip file containing test vectors for this set of algorithms
  url: 'https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Algorithm-Validation-Program/documents/sha3/shakebytetestvectors.zip',

  # expected sha256 digest of zip file
  hash: 'debfebc3157b3ceea002b84ca38476420389a3bf7e97dc5f53ea4689a16de4c7',

  # algorithms
  algos: [{
    name: 'SHAKE128', # algo name
    fn: 'shake128', # function prefix
    type: :xof, # function type

    # response files
    rsp_files: %w{
      SHAKE128ShortMsg.rsp
      SHAKE128LongMsg.rsp
      SHAKE128VariableOut.rsp
    },
  }, {
    name: 'SHAKE256', # algo name
    fn: 'shake256', # function prefix
    type: :xof, # function type

    # response files
    rsp_files: %w{
      SHAKE256ShortMsg.rsp
      SHAKE256LongMsg.rsp
      SHAKE256VariableOut.rsp
    },
  }],
}]

# main.c header
MAIN_HEAD = <<END_MAIN_HEAD
// Run test vectors for FIPS202 hash functions and XOFs.
// (generated by `gen-main.rb`)

#include <stdint.h> // uint8_t
#include <stdio.h> // printf()
#include <string.h> // memcmp()
#include "hex.h" // hex_write()
#include "sha3.h"

// test failure
typedef struct {
  const char *test_fn; // test function name
  size_t test_i; // test case index
  const char *rsp_file; // response file name
  size_t rsp_line; // start line of test in response file
  const uint8_t *got; // result
  size_t got_len; // result length
  const uint8_t *exp; // expected result
  size_t exp_len; // expected result length
} failure_t;

// print failure
static void fail(FILE *fh, const failure_t f) {
  fprintf(fh, "test_fn = %s\\n", f.test_fn);
  fprintf(fh, "rsp_file = %s\\n", f.rsp_file);
  fprintf(fh, "rsp_line = %zu\\n", f.rsp_line);
  fprintf(fh, "got (%zu) = ", f.got_len);
  hex_write(fh, f.got, f.got_len);
  fprintf(fh, "\\nexp (%zu) = ", f.exp_len);
  hex_write(fh, f.exp, f.exp_len);
  fprintf(fh, "\\n");
}
END_MAIN_HEAD

# main.c footer template
MAIN_TAIL_TMPL = <<END_MAIN_TAIL_TMPL
int main(void) {
%<fns>s
  return 0;
}
END_MAIN_TAIL_TMPL

# xof test vector template
XOF_TEST_TMPL = <<END_XOF_TEST_TMPL
    { %<file_pos>d, %<line>d, %<src_ofs>d, %<src_len>d, %<exp_ofs>d, %<exp_len>d },
END_XOF_TEST_TMPL

# xof test function template
XOF_FN_TMPL = <<END_XOF_FN_TMPL
// %<name>s tests from shakebytetestvectors.zip
static void test_%<fn>s(void) {
  // data for tests
  static const uint8_t DATA[] = {
%<data>s
  }; // total size = %<data_size>d

  // response file names
  static const char *RSP_FILES[] = { %<rsp_files>s };

  // test vectors
  static const struct {
    const size_t rsp_file_ofs, // offset of source file name in RSP_FILES
                 rsp_file_line, // line within rsp file
                 src_ofs, // message offset into DATA
                 src_len, // message length of data
                 exp_ofs, // expected output offset into DATA
                 exp_len; // expected output length
  } TESTS[] = {
%<tests>s
  };

  // run tests
  for (size_t i = 0; i < sizeof(TESTS) / sizeof(TESTS[0]); i++) {
    // get expected result
    const uint8_t * const exp = DATA + TESTS[i].exp_ofs;
    const size_t exp_len = TESTS[i].exp_len;

    // hash data into "got"
    uint8_t got[%<max_exp_len>d] = { 0 };
    %<fn>s(DATA + TESTS[i].src_ofs, TESTS[i].src_len, got, exp_len);

    // check for expected result
    if (memcmp(got, exp, exp_len)) {
      fail(stderr, (failure_t) {
        .test_fn = __func__, // test function
        .test_i = i, // test case index in TESTS
        .rsp_file = RSP_FILES[TESTS[i].rsp_file_ofs], // response file
        .rsp_line = TESTS[i].rsp_file_line, // start line of test case
        .got = got, // result
        .got_len = exp_len, // result length, in bytes
        .exp = exp, // expected result
        .exp_len = exp_len, // expected result length, in bytes
      });
    }
  }
}
END_XOF_FN_TMPL

# hash test vector template
HASH_TEST_TMPL = <<END_XOF_TEST_TMPL
    { %<file_pos>d, %<line>d, %<src_ofs>d, %<src_len>d, %<exp_ofs>d },
END_XOF_TEST_TMPL

# hash test function template
HASH_FN_TMPL = <<END_XOF_FN_TMPL
// %<name>s tests
static void test_%<fn>s(void) {
  // data for tests
  static const uint8_t DATA[] = {
%<data>s
  }; // total size = %<data_size>d

  // response file names
  static const char *RSP_FILES[] = { %<rsp_files>s };

  // test vectors
  static const struct {
    const size_t rsp_file_ofs, // offset of source file name in RSP_FILES
                 rsp_file_line, // line within rsp file
                 src_ofs, // message offset into DATA
                 src_len, // message length of data
                 exp_ofs; // expected output offset into DATA
  } TESTS[] = {
%<tests>s
  };

  // run tests
  for (size_t i = 0; i < sizeof(TESTS) / sizeof(TESTS[0]); i++) {
    // get expected result
    const uint8_t * const exp = DATA + TESTS[i].exp_ofs;

    // hash data into "got"
    uint8_t got[%<max_exp_len>d] = { 0 };
    %<fn>s(DATA + TESTS[i].src_ofs, TESTS[i].src_len, got);

    // check for expected result
    if (memcmp(got, exp, sizeof(got))) {
      fail(stderr, (failure_t) {
        .test_fn = __func__, // test function
        .test_i = i, // test case index in TESTS
        .rsp_file = RSP_FILES[TESTS[i].rsp_file_ofs], // response file
        .rsp_line = TESTS[i].rsp_file_line, // start line of test case
        .got = got, // result
        .got_len = sizeof(got), // result length, in bytes
        .exp = exp, // expected result
        .exp_len = sizeof(got), // expected result length, in bytes
      });
    }
  }
}
END_XOF_FN_TMPL

# parse hash of key/value pairs into test vector row
def parse(algo_type, file, h, data_size, &block)
  # parse hex-encoded message into array of bytes, then
  # truncate message to "Len" bytes, if "Len" is specified
  src = h[:Msg].scan(/../).map { |s| s.to_i(16) }
  src = src[0, h[:Len].to_i] if h.key?(:Len)
  src_ofs = data_size

  exp = case algo_type
  when :xof
    # parse hex-encoded expected output into array of bytes, then
    # truncate expected output to "OutputLen" bytes, if "OutputLen" is
    # specified
    exp = h[:Output].scan(/../).map { |s| s.to_i(16) }
    exp = exp[0, h[:OutputLen].to_i] if h.key?(:OutputLen)
    exp
  when :hash
    # get expected output
    h[:MD].scan(/../).map { |s| s.to_i(16) }
  else
    raise "unknown algorithm type: #{algo_type}"
  end
  exp_ofs = data_size + src.size

  # append message bytes and expected output bytes to data
  block.call(src + exp)

  # return parsed entry
  {
    file: file,
    line: h[:line],

    # message
    src_ofs: src_ofs,
    src_len: src.size,

    # expected output
    exp_ofs: exp_ofs,
    exp_len: exp.size,
  }
end

# build path to test vector archive cache
CACHE_DIR = File.join(__dir__, 'zip-cache')
Dir.mkdir(CACHE_DIR) unless Dir.exist?(CACHE_DIR)

# download test vector archives
SETS.map do |set|
  Thread.new(set) do |set|
    # parse url
    url = URI.parse(set[:url])

    # destination archive name
    zip_name = File.basename(url.path)

    # absolute path to destination archive
    zip_path = File.join(CACHE_DIR, zip_name)

    # download to destination path
    url.open { |src_io| IO.copy_stream(src_io, zip_path) } unless File.exist?(zip_path)

    # get sha256 hash of downloaded archive
    got_hash = ::OpenSSL::Digest::SHA256.new.file(zip_path).hexdigest

    # check against expected digest
    if got_hash != set[:hash]
      raise '%s: hash mismatch: got %s, exp %s' % [zip_name, got_hash, set[:hash]]
    end
  end
end.each { |t| t.join }

# print header
puts(MAIN_HEAD)

# generate test code
SETS.each do |set|
  # get absolute path to zip file
  zip_path = File.join(CACHE_DIR, File.basename(URI.parse(set[:url]).path))

  # open zip file
  Zip::File.open(zip_path, 'rb') do |zip|
    set[:algos].each do |algo|
      data = [] # test data (shared across all rsp files)

      # parse test vectors from rsp files
      rows = algo[:rsp_files].each_with_object([]) do |rsp_file, rows|
        curr = {} # current row

        # read lines from rsp file
        lines = zip.glob(rsp_file).first.get_input_stream.readlines.to_a.map { |line| line.strip }

        # parse lines into rows
        lines.size.times.each do |line_i|
          case lines[line_i]
          when /^(\w+) = (\w+)$/
            k, v = $1, $2 # extract key and value
            curr[k.intern] = v

            # cache line number
            curr[:line] = line_i unless curr[:line]
          when ''
            if curr.size > 0
              # parse current hash into test vector and add it to results
              rows << parse(algo[:type], rsp_file, curr, data.size) { |bytes| data += bytes }
            end

            curr = {} # clear current hash
          end
        end
      end

      # get maximum number of lines of emitted static DATA array
      num_lines = data.size / PER_LINE + ((data.size % PER_LINE) > 0 ? 1 : 0)

      # get maximum length of expected output
      # used to for size of "got" buffer in emitted code
      max_exp_len = case algo[:type]
      when :hash
        algo[:size]
      when :xof
        rows.reduce(0) { |r, row| row[:exp_len] > r ? row[:exp_len] : r }
      else
        raise "unknown algorithm type: #{algo[:type]}"
      end

      # get templates
      fn_tmpl, test_tmpl = case algo[:type]
      when :xof
        [XOF_FN_TMPL, XOF_TEST_TMPL]
      when :hash
        [HASH_FN_TMPL, HASH_TEST_TMPL]
      else
        raise "unknown algorithm type: #{algo[:type]}"
      end

      # generate test function
      puts(fn_tmpl % {
        name: algo[:name], # algorithm name
        fn: algo[:fn], # algorithm function prefix

        # test vector data
        data: num_lines.times.map { |ofs|
          "    %s,\n" % [data[PER_LINE * ofs, PER_LINE].join(', ')]
        }.join,

        # response file names
        rsp_files: algo[:rsp_files].map { |file| '"%s"' % [file] }.join(', '),

        data_size: data.size, # total size of emitted DATA array, in bytes
        max_exp_len: max_exp_len, # maximum expected output length, in bytes

        # test vectors
        tests: rows.map { |row|
          test_tmpl % row.merge({ file_pos: algo[:rsp_files].index(row[:file]) })
        }.join,
      })
    end
  end
end

# print main function body
puts(MAIN_TAIL_TMPL % {
  fns: SETS.each_with_object([]) do |set, r|
    set[:algos].each_with_object(r) do |algo, r|
      r << '  test_%<fn>s();' % algo
    end
  end.join("\n"),
})