/* Copyright (c) 2021 Contrast Security, Inc.  See
 * https://www.contrastsecurity.com/enduser-terms-0317a for more details. */

#include "cs__assess_module.h"
#include "../cs__common/cs__common.h"
#include <ruby.h>

void contrast_assess_eval_trigger_check(VALUE module, VALUE source, VALUE ret) {
    if (RTEST(
            rb_funcall(contrast_patcher(), rb_sym_skip_contrast_analysis, 0))) {
        return;
    }

    int nested_scope =
        RTEST(rb_funcall(contrast_patcher(), rb_sym_in_scope, 0));

    rb_funcall(contrast_patcher(), rb_sym_enter_scope, 0);

    if (!nested_scope) {
        VALUE method = rb_funcall(rb_mKernel, rb_sym_method, 0);
        /* If this method ever throws an exception, the scope-leave
         * needs to be moved within a rescue call.
         */
        rb_funcall(module_eval_trigger, trigger_check_method, 4, module, source,
                   ret, method);
    }

    rb_funcall(contrast_patcher(), rb_sym_exit_scope, 0);
}

VALUE
contrast_assess_module_class_eval(const int argc, const VALUE *argv,
                                  const VALUE mod) {
    VALUE ret = rb_mod_module_eval(argc, argv, mod);

    if (argc > 0) {
        VALUE data = argv[0];
        contrast_assess_eval_trigger_check(mod, data, ret);
    }
    rb_funcall(assess_patcher, rb_sym_assess_patch_eval, 1, mod);

    return ret;
}

VALUE
contrast_assess_module_module_eval(const int argc, const VALUE *argv,
                                   const VALUE mod) {
    VALUE ret = rb_mod_module_eval(argc, argv, mod);

    if (argc > 0) {
        VALUE data = argv[0];
        contrast_assess_eval_trigger_check(mod, data, ret);
    }

    rb_funcall(assess_patcher, rb_sym_assess_patch_eval, 1, mod);

    return ret;
}

VALUE
contrast_assess_module_prepend(const int argc, const VALUE *argv,
                               const VALUE self) {

    rb_prepend_module(self, argv[0]);

    VALUE module_at;
    VALUE rb_incl_in_mod_ary = rb_funcall(self, rb_intern("included_in"), 0);

    if (RB_TYPE_P(rb_incl_in_mod_ary, T_ARRAY)) {
       int i = 0;
       int size = rb_funcall(rb_incl_in_mod_ary, rb_intern("length"), 0);
       for (i = 0; i < size; ++i) {
           module_at = rb_ary_entry(rb_incl_in_mod_ary, i);
           if (RB_TYPE_P(module_at, T_MODULE)) {
              rb_include_module(module_at, argv[0]);
           }
       }
    }
    return self;
}

VALUE
contrast_assess_module_included(const int argc, const VALUE *argv,
                               const VALUE self) {
   VALUE frozen;
   if (RB_TYPE_P(self, T_MODULE)) {
       // check if frozen
       frozen = rb_funcall(self, rb_intern("cs__frozen?"), 0);
       if (frozen == Qfalse) {
           VALUE ary = rb_funcall(self, rb_intern("included_in"), 0);
           if (RB_TYPE_P(ary, T_ARRAY)) {
              rb_ary_push(ary, argv[0]);
           }
       }
   }
   return self;
}

void Init_cs__assess_module(void) {
    module_eval_trigger =
        rb_define_class_under(core_assess, "EvalTrigger", rb_cObject);
    trigger_check_method = rb_intern("eval_trigger_check");

    rb_sym_assess_patch_eval = rb_intern("patch_assess_on_eval");
    assess_patcher = rb_define_module_under(assess_policy, "Patcher");

    /* Returns of these 2 patches are discarded.
     * We're calling the underlying via direct C, instead of
     * whatever method was there before.
     * See similar comments in basic_object C ext patch.
     */

    contrast_register_patch("Module", "class_eval",
                            contrast_assess_module_class_eval);

    contrast_register_patch("Module", "module_eval",
                            contrast_assess_module_module_eval);
   /*
    * We patch these for better ancestors handling, and only for older ruby versions.
    */
  if (rb_ver_below_three()) {
     contrast_register_patch("Module", "included",
                             contrast_assess_module_included);
     contrast_register_patch("Module", "prepend",
                             contrast_assess_module_prepend);
     }
}