#! /usr/bin/env ruby
#

require 'getoptlong'
require 'tmpdir'
require 'set'

$random_length = 32
$prefix = "stress"
$max_rules = 50
$min_rules = 5

def get_random_name(len=$random_length)
  return sprintf("%0#{len}x", rand(2 ** (4 * len)))
end

def get_random_regex()
  case rand(10)
    when 0..3
      return "{#{get_random_name(rand(8) + 2)},#{get_random_name(rand(8) + 2)},#{get_random_name(rand(8) + 2)}}"
    when 4..5
      return "[#{get_random_name(rand(5) + 1)}]"
    when 6..7
      return "*"
    when 8..9
      return "**"
  end
end

def get_random_path(symtab)
  # Always prefix with a non-regex element
  out = "/#{get_random_name(rand(10) + 4)}"
  0.upto(rand(20) + 2) do
    case rand(50)
      when 0
        out = "#{out}/@{#{symtab.get_symbol}}"
      when 1..4
        out = "#{out}/#{get_random_regex}"
      when 5..49
        out = "#{out}/#{get_random_name(rand(10) + 4)}"
    end
  end
  return out
end

def get_random_mode()
  case rand(10)
    when 0..4
      return "r"
    when 5..7
      return "rw"
    when 8
      return "Px"
    when 9
      return "rix"
  end
end

class SymTab
  include Enumerable

  def initialize()
    @symtab = { }
  end

  def dump
    @symtab.each do | symbol, values |
      out = "@{#{symbol}}="
      values.each { |v| out += " #{v}"}
      puts out
    end
  end

  def each_decl_s
    @symtab.each do | symbol, values |
      out = "@{#{symbol}}="
      values.each { |v| out += " #{v}"}
      yield(out)
    end
  end

  def get_symbol
    # variables need to be prefixed with letters, hence the VAR prefix
    symbol = "VAR_" + get_random_name(rand(6) + 4)
    values = []
    0.upto(rand(4)) { values << get_random_name(rand(8) + 2) }
    @symtab[symbol] = values
    return symbol
  end
end

# Abstract class, though may become a real class for generation of
# random types of rules
class Rule
end

class FileRule < Rule
  def initialize(symtab=nil)
    @path = get_random_path(symtab)
    @mode = get_random_mode()
  end

  def to_s
    return "#{@path} #{@mode},"
  end
end

class NamedFileRule < FileRule
  def initialize(path, mode)
    @path = path
    @mode = mode
  end
end

class CapRule < Rule
  CAP_LIST = [
    "chown",
    "dac_override",
    "dac_read_search",
    "fowner",
    "fsetid",
    "kill",
    "setgid",
    "setuid",
    "setpcap",
    "linux_immutable",
    "net_bind_service",
    "net_broadcast",
    "net_admin",
    "net_raw",
    "ipc_lock",
    "ipc_owner",
    "sys_module",
    "sys_rawio",
    "sys_chroot",
    "sys_ptrace",
    "sys_pacct",
    "sys_admin",
    "sys_boot",
    "sys_nice",
    "sys_resource",
    "sys_time",
    "sys_tty_config",
    "mknod",
    "lease",
    "audit_write",
    "audit_control",
    "setfcap",
    "mac_override",
    "mac_admin"
  ]

  def initialize()
    @cap = CAP_LIST[rand(CAP_LIST.length)]
  end

  def to_s
    return "capability #{@cap},"
  end
end

class NetRule < Rule
  # XXX Fill me in
end

class RlimitRule < Rule
  RLIMIT_LIST = [
#"cpu",		# cpu rlimit not supported
    "fsize",
    "data",
    "stack",
    "core",
    "rss",
    "nofile",
    "ofile",
    "as",
    "nproc",
    "memlock",
    "locks",
    "sigpending",
    "msgqueue",
    "nice",
    "rtprio"
  ]

  def initialize()
    @rlimit = RLIMIT_LIST[rand(RLIMIT_LIST.length)]
    if rand(20) == 0
      @limit = "infinity"
    elsif @rlimit == "nice"
      @limit = rand(40) - 20
    else
      @limit = rand(2 ** 31)
    end
  end

  def to_s
    return "set rlimit #{@rlimit} <= #{@limit},"
  end
end

class Flags
  FLAG_LIST = [
    "complain",
    "audit",
    "chroot_relative",
    "namespace_relative",
    "mediate_deleted",
    "delegate_deleted",
    "attach_disconnected",
    "no_attach_disconnected",
    "chroot_attach",
    "chroot_no_attach"
  ]

  FLAG_CONFLICTS = [
    ["chroot_relative", "namespace_relative"],
    ["mediate_deleted", "delegate_deleted"],
    ["attach_disconnected", "no_attach_disconnected"],
    ["chroot_attach", "chroot_no_attach"]
  ]

  def initialize()
    @flags = Set.new()
    if rand(2) == 1
      return
    end

    0.upto(4 - Math.log(rand(32) + 1).to_int) do |x|
      @flags << FLAG_LIST[rand(FLAG_LIST.length)]
    end

    FLAG_CONFLICTS.each do |c|
      if @flags.include?(c[0]) and @flags.include?(c[1])
        @flags.delete(c[rand(2)])
      end
    end
  end

  def to_s
    if @flags.empty?
      return ""
    end
    out = @flags.to_a.join(",")
    return "flags=(#{out})"
  end
end

def prefix_to_s(name)
  out = []
  out << "#"
  out << "# prefix for #{name}"
  out << "# generated by #{__FILE__}"
  out << "#include <tunables/global>"
  out << "#"
end

class Profile
  attr_reader :rvalue
  attr_reader :name

  def initialize()
    @rvalue = get_random_name()
    @name = "/does/not/exist/#{@rvalue}"
    @rules = []
    @flags = Flags.new()
    @symtab = SymTab.new()
  end

  def generate_rules
    @rules << NamedFileRule.new(@name, "rm").to_s
    0.upto(rand($max_rules - $min_rules) + $min_rules) do |x|
      case rand(100)
        when 0..14
          @rules << CapRule.new.to_s
        when 15..24
          @rules << RlimitRule.new.to_s
        when 25..100
          @rules << FileRule.new(symtab=@symtab).to_s
      end
    end
  end

  def decl_s
    out = []
    out << "#"
    out << "# variable declarations for #{@name}"
    out << "# generated by #{__FILE__}"
    out << "#"
    @symtab.each_decl_s { | decl | out << decl }
    out << ""
  end

  def to_s
    out = []
    out << "#"
    out << "# profile for #{@name}"
    out << "# generated by #{__FILE__}"
    out << "#"
    out << "#{@name} #{@flags} {"
    out << "  #include <abstractions/base>"
    out << ""
    @rules.sort.each { |r| out << "  #{r}" }
    out << "}"
    out << ""
  end
end

def showUsage
  warn "usage: #{__FILE__} count"
  exit 1
end

def gen_profiles_dir(profiles)
  # Mu, no secure tmpdir creation in base ruby
  begin
    dirname = "#{Dir.tmpdir}/#{$prefix}-#{get_random_name(32)}"
    Dir.mkdir(dirname, 0755)
  rescue Errno::EEXIST
    retry
  end

  profiles.each do |p|
    open("#{dirname}/#{p.rvalue}.sd", File::CREAT|File::EXCL|File::WRONLY, 0644) do |file|
      file.puts(prefix_to_s(p.name))
      file.puts(p.decl_s)
      file.puts(p.to_s)
    end
  end

  return dirname
end

def gen_profiles_file(profiles)
  # Mu, no secure tempfile creation in base ruby
  begin
    filename = "#{Dir.tmpdir}/#{$prefix}-#{get_random_name(32)}.sd"
    File.open(filename, File::CREAT|File::EXCL|File::WRONLY, 0644) do |file|
      file.puts(prefix_to_s(filename))
      profiles.each { |p| file.puts(p.decl_s) }
      profiles.each { |p| file.puts(p.to_s) }
    end
  rescue Errno::EEXIST
    retry
  end

  return filename
end

if __FILE__ == $0

  keep_files = true
  opts = GetoptLong.new(
    [ '--help', '-h', GetoptLong::NO_ARGUMENT ],
    [ '--seed', '-s', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--keep-files', '-k', GetoptLong::NO_ARGUMENT ],
    [ '--max-rules', '-M', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--min-rules', '-m', GetoptLong::REQUIRED_ARGUMENT ]
  )

  opts.each do |opt, arg|
    case opt
      when '--help'
	showUsage
      when '--seed'
        srand(arg.to_i)
      when '--keep-files'
        keep_files = true
      when '--max-rules'
        $max_rules = arg.to_i
      when '--min-rules'
        $min_rules = arg.to_i
    end
  end

  showUsage if ARGV.length != 1
  count = ARGV.shift.to_i
  showUsage if count < 1
  profiles = []
  while profiles.length < count do
    profiles << Profile.new()
  end
  profiles.each { |p| p.generate_rules }

  begin
    profiles_dir = gen_profiles_dir(profiles)
    profile_single = gen_profiles_file(profiles)

    ensure
      if (keep_files == false)
        Dir.foreach(profiles_dir) do |filename|
          File.delete("#{profiles_dir}/#{filename}") if (filename != '.' and filename != '..')
        end
        Dir.rmdir(profiles_dir)
  	File.delete(profile_single)
      else
        puts "PROFILEDIR=#{profiles_dir}; export PROFILEDIR"
	puts "PROFILESINGLE=#{profile_single}; export PROFILESINGLE"
      end
  end
end
