#!/usr/bin/env ruby
###################################################
# ----------------------------------------------- #
# Fuzz Origami Ruby gem with mutated PDF files    #
# ----------------------------------------------- #
#                                                 #
# Each test case is written to 'fuzz.pdf' in the  #
# current working directory.                      #
#                                                 #
# Crashes and the associated backtrace are saved  #
# in the 'crashes' directory in the current       #
# working directory.                              #
#                                                 #
###################################################
# ~ bcoles

require 'date'
require 'origami'
require 'colorize'
require 'fileutils'
require 'timeout'
require 'securerandom'

include Origami

VERBOSE = false
OUTPUT_DIR = "#{Dir.pwd}/crashes".freeze

#
# Show usage
#
def usage
  puts 'Usage: ./fuzz.rb <FILE1> [FILE2] [FILE3] [...]'
  puts 'Example: ./fuzz.rb test/dataset/**.pdf'
  exit 1
end

#
# Print status message
#
# @param [String] msg message to print
#
def print_status(msg = '')
  puts '[*] '.blue + msg if VERBOSE
end

#
# Print progress messages
#
# @param [String] msg message to print
#
def print_good(msg = '')
  puts '[+] '.green + msg if VERBOSE
end

#
# Print error message
#
# @param [String] msg message to print
#
def print_error(msg = '')
  puts '[-] '.red + msg
end

#
# Setup environment
#
def setup
  FileUtils.mkdir_p OUTPUT_DIR unless File.directory? OUTPUT_DIR
rescue => e
  print_error "Could not create output directory '#{OUTPUT_DIR}': #{e}"
  exit 1
end

#
# Generate a mutated PDF file with a single mitated byte
#
# @param [Path] f path to PDF file
#
def mutate_byte(f)
  data = IO.binread f
  position = SecureRandom.random_number data.size
  new_byte = SecureRandom.random_number 256
  new_data = data.dup.tap { |s| s.setbyte(position, new_byte) }

  File.open(@fuzz_outfile, 'w') do |file|
    file.write new_data
  end
end

#
# Generate a mutated PDF file with multiple mutated bytes
#
# @param [Path] f path to PDF file
#
def mutate_bytes(f)
  data = IO.binread f
  fuzz_factor = 200
  num_writes = rand((data.size / fuzz_factor.to_f).ceil) + 1

  new_data = data.dup
  num_writes.times do
    position = SecureRandom.random_number data.size
    new_byte = SecureRandom.random_number 256
    new_data.tap { |stream| stream.setbyte position, new_byte }
  end

  File.open(@fuzz_outfile, 'w') do |file|
    file.write new_data
  end
end

#
# Generate a mutated PDF file with all integers replaced by '-1'
#
# @param [Path] f path to PDF file
#
def clobber_integers(f)
  data = IO.binread f
  new_data = data.dup.gsub(/\d/, '-1')

  File.open(@fuzz_outfile, 'w') do |file|
    file.write new_data
  end
end

#
# Generate a mutated PDF file with all strings 3 characters or longer
# replaced with 2000 'A' characters
#
# @param [Path] f path to PDF file
#
def clobber_strings(f)
  data = IO.binread f
  new_data = data.dup.gsub(/[a-zA-Z]{3,}/, 'A' * 2000)

  File.open(@fuzz_outfile, 'w') do |file|
    file.write new_data
  end
end

#
# Read a PDF file
#
# @param [String] f path to PDF file
#
def read(f)
  print_status "Processing '#{f}'"
  begin
    pdf = PDF.read f, ignore_errors: false, password: 'invalid'
  rescue Origami => e
    print_status "Could not read PDF '#{f}': PDF is malformed"
    return
  end
  print_good 'Processing complete'

  print_status "Parsing '#{f}'"
  begin
    parse pdf
  rescue Origami => e
    print_status "Could not parse PDF '#{f}': PDF is malformed"
    return
  end
  print_good 'Parsing complete'
end

#
# Parse PDF
#
def parse(pdf)
  print_status "Version: #{pdf.header.to_f}"
  print_status "Number of revisions: #{pdf.revisions.size}"
  print_status "Number of indirect objects: #{pdf.indirect_objects.size}"
  print_status "Number of pages: #{pdf.pages.count}"
  print_status "Linearized: #{pdf.linearized?}"
  print_status "Encrypted: #{pdf.encrypted?}"
  print_status "Signed: #{pdf.signed?}"
  print_status "Has usage rights: #{pdf.usage_rights?}"
  print_status "Form: #{pdf.form?}"
  print_status "XFA form: #{pdf.xfa_form?}"
  print_status "Document Info: #{pdf.document_info?}"
  print_status "Metadata: #{pdf.metadata?}"

  print_status 'Parsing PDF contents...'
  pdf.pages.each_with_index do |page, index|
    page.each_annotation do |a|
    end
    page.each_content_stream do |a|
    end
    page.each_annotation do |a|
    end
    page.each_colorspace do |a|
    end
    page.each_extgstate do |a|
    end
    page.each_pattern do |a|
    end
    page.each_shading do |a|
    end
    page.each_xobject do |a|
    end
    page.each_font do |a|
    end
    page.each_property do |a|
    end
    page.each do |a|
    end
    page.each_key do |a|
    end
    page.each_value do |a|
    end
    page.each_pair do |a|
    end
    page.each_with_index do |a, b|
    end
    page.reverse_each do |a|
    end
    page.each_entry do |a|
    end
    #page.each_resource do |a|
    #end
    #page.each_slice do |a|
    #end
    #page.each_cons do |a|
    #end
    #page.each_with_object do |a|
    #end
  end
end

#
# Show summary of crashes
#
def summary
  puts
  puts "Complete! Crashes saved to '#{OUTPUT_DIR}'"
  puts
  puts `/usr/bin/head -n1 #{OUTPUT_DIR}/*.trace` if File.exist? '/usr/bin/head'
end

#
# Report error message to STDOUT
# and save fuzz test case and backtrace to OUTPUT_DIR
#
def report_crash(e)
  puts " - #{e.message}"
  puts e.backtrace.first
  fname = "#{DateTime.now.strftime('%Y%m%d%H%M%S%N')}_crash_#{rand(1000)}"
  FileUtils.mv @fuzz_outfile, "#{OUTPUT_DIR}/#{fname}.pdf"
  File.open("#{OUTPUT_DIR}/#{fname}.pdf.trace", 'w') do |file|
    file.write "#{e.message}\n#{e.backtrace.join "\n"}"
  end
end

#
# Test Origami with the mutated file
#
def test
  Timeout.timeout(@timeout) do
    read @fuzz_outfile
  end
rescue SystemStackError => e
  report_crash e
rescue Timeout::Error => e
  report_crash e
rescue SyntaxError => e
  report_crash e
rescue => e
  raise e unless e.backtrace.join("\n") =~ %r{gems/origami}
  report_crash e
end

#
# Generate random byte mutations and run test
#
# @param [String] f path to PDF file
#
def fuzz_bytes(f)
  iterations = 1000
  1.upto(iterations) do |i|
    print "\r#{(i * 100) / iterations} % (#{i} / #{iterations})"
    mutate_bytes f
    test
  end
end

#
# Generate integer mutations and run tests
#
# @param [String] f path to PDF file
#
def fuzz_integers(f)
  clobber_integers f
  test
end

#
# Generate string mutations and run tests
#
# @param [String] f path to PDF file
#
def fuzz_strings(f)
  clobber_strings f
  test
end

puts '-' * 60
puts '% Fuzzer for Origami Ruby gem'
puts '-' * 60
puts

usage if ARGV[0].nil?

setup

@timeout = 15
@fuzz_outfile = 'fuzz.pdf'

trap 'SIGINT' do
  puts
  puts 'Caught interrupt. Exiting...'
  summary
  exit 130
end

ARGV.each do |f|
  unless File.exist? f
    print_error "Could not find file '#{f}'"
    next
  end

  fuzz_integers f
  fuzz_strings f
  fuzz_bytes f

  puts '-' * 60
end

summary