/*
 * pi3's Linux kernel Runtime Guard
 *
 * Component:
 *  - x86/AMD64 specific code for Exploit Detection
 *
 * Notes:
 *  - None
 *
 * Timeline:
 *  - Created: 18.IV.2019
 *
 * Author:
 *  - Adam 'pi3' Zabrocki (http://pi3.com.pl)
 *
 */

#ifdef CONFIG_X86

#ifndef P_LKRG_EXPLOIT_DETECTION_X86_ARCH_H
#define P_LKRG_EXPLOIT_DETECTION_X86_ARCH_H

/*
 * pCFI global CPU flags
 */
#define P_PCFI_X86_WP     0x1
#define P_PCFI_X86_SMEP   0x2
#define P_PCFI_X86_SMAP   0x4

#define P_IS_WP_FLAG_ENABLED(x)   (x & P_PCFI_X86_WP)
#define P_IS_SMEP_FLAG_ENABLED(x) (x & P_PCFI_X86_SMEP)
#define P_IS_SMAP_FLAG_ENABLED(x) (x & P_PCFI_X86_SMAP)

/*
 * Known minor bug: these updates are non-atomic, so concurrent changing of
 * these flags might result in all but one of the changes getting lost.  We
 * may switch to using individual Boolean variables instead of the bitmask.
 */
#define P_ENABLE_WP_FLAG(x)       (x |= P_PCFI_X86_WP)
#define P_ENABLE_SMEP_FLAG(x)     (x |= P_PCFI_X86_SMEP)
#define P_ENABLE_SMAP_FLAG(x)     (x |= P_PCFI_X86_SMAP)

#define P_DISABLE_WP_FLAG(x)      (x &= ~P_PCFI_X86_WP)
#define P_DISABLE_SMEP_FLAG(x)    (x &= ~P_PCFI_X86_SMEP)
#define P_DISABLE_SMAP_FLAG(x)    (x &= ~P_PCFI_X86_SMAP)

/*
 * x86 specific functions
 */
static inline void p_write_cr4(unsigned long p_arg) {

#if LINUX_VERSION_CODE >= KERNEL_VERSION(5,8,0)
   P_SYM(p_native_write_cr4)(p_arg);
#elif LINUX_VERSION_CODE < KERNEL_VERSION(4,0,0)
   write_cr4(p_arg);
#else
   __write_cr4(p_arg);
#endif

}

static inline unsigned long p_read_cr4(void) {

#if LINUX_VERSION_CODE < KERNEL_VERSION(4,0,0)
   return read_cr4();
#else
   return __read_cr4();
#endif

}

static inline void p_kill_current(void) {
   /*
    * It is possible to kill an almost arbitrary process here. Such situation
    * is possible if process which disables WP/SMEP/SMAP gets preempted before
    * our detection. Most of the exploits which I tested didn't have such FP,
    * but such possibility of course exists. However, it is worth the risk.
    * Especially given that such changes are not legal and are always a result
    * of malicious activity.
    */
   struct task_struct *p_task = current;
   if (P_CTRL(p_pint_enforce) == 1 && task_pid_nr(p_task) != 1) {
      p_print_log(P_LOG_ALERT, "BLOCK: Task: About to take chances and kill a task that may be relevant or not");
      p_kill_task_by_task(p_task);
   }
}

static inline unsigned int p_ed_pcfi_x86_validate_wp(unsigned char p_kill) {

   unsigned int p_bad = 0;
   register unsigned long p_cr0;

   if (!P_CTRL(p_kint_validate))
      return p_bad;

   p_cr0 = read_cr0();

   if (!(p_cr0 & X86_CR0_WP)) {

      switch (P_CTRL(p_kint_enforce)) {

         /* Panic */
         case 2:
            // OK, we need to crash the kernel now
            p_panic("CPU: Write Protection got disabled");
            break;

         /* Log and restore */
         case 1:
            p_print_log(P_LOG_ALERT, "BLOCK: CPU: Write Protection got disabled, re-enabling");
            p_cr0 |= X86_CR0_WP;
            write_cr0(p_cr0);
            wbinvd();
            if (p_kill)
               p_kill_current();
            p_bad++;
            break;

         /* Log and accept */
         case 0:
            if (!P_IS_WP_FLAG_ENABLED(p_pcfi_CPU_flags))
               return p_bad;
            p_print_log(P_LOG_ALERT, "ALLOW: CPU: Write Protection got disabled, accepting");
            P_DISABLE_WP_FLAG(p_pcfi_CPU_flags);
            break;
      }
   }

   return p_bad;
}

static inline unsigned int p_ed_pcfi_x86_validate_smXp(unsigned char p_kill) {

   unsigned int p_bad = 0;
   register unsigned long p_cr4;


   if (!P_CTRL(p_smep_validate) && !P_CTRL(p_smap_validate))
      return p_bad;

   p_cr4 = p_read_cr4();

   if (!(p_cr4 & X86_CR4_SMEP) && P_CTRL(p_smep_validate)) {

      switch (P_CTRL(p_smep_enforce)) {

         /* Panic */
         case 2:
            // OK, we need to crash the kernel now
            p_panic("CPU: SMEP got disabled");
            break;

         /* Log and restore */
         case 1:
            p_print_log(P_LOG_ALERT, "BLOCK: CPU: SMEP got disabled, re-enabling");
//            cr4_set_bits(X86_CR4_SMEP);
            p_cr4 |= X86_CR4_SMEP;
            p_write_cr4(p_cr4);
            if (p_kill)
               p_kill_current();
            p_bad++;
            break;

         /* Log and accept */
         case 0:
            if (!P_IS_SMEP_FLAG_ENABLED(p_pcfi_CPU_flags))
               return p_bad;
            p_print_log(P_LOG_ALERT, "ALLOW: CPU: SMEP got disabled, accepting");
            P_DISABLE_SMEP_FLAG(p_pcfi_CPU_flags);
            break;
      }
   }

   if (!(p_cr4 & X86_CR4_SMAP) && P_CTRL(p_smap_validate)) {

      switch (P_CTRL(p_smap_enforce)) {

         /* Panic */
         case 2:
            // OK, we need to crash the kernel now
            p_panic("CPU: SMAP got disabled");
            break;

         /* Log and restore */
         case 1:
            p_print_log(P_LOG_ALERT, "BLOCK: CPU: SMAP got disabled, re-enabling");
//            cr4_set_bits(X86_CR4_SMAP);
            p_cr4 |= X86_CR4_SMAP;
            p_write_cr4(p_cr4);
            if (p_kill)
               p_kill_current();
            p_bad++;
            break;

         /* Log and accept */
         case 0:
            if (!P_IS_SMAP_FLAG_ENABLED(p_pcfi_CPU_flags))
               return p_bad;
            p_print_log(P_LOG_ALERT, "ALLOW: CPU: SMAP got disabled, accepting");
            P_DISABLE_SMAP_FLAG(p_pcfi_CPU_flags);
            break;
      }
   }

   return p_bad;
}

#endif

#endif
