#define STP_NETLINK_ONLY
#define STP_NUM_STRINGS 1

#include "runtime.h"

#define NEED_INT64_VALS
#define NEED_STRING_VALS

#define KEY1_TYPE INT64
#include "map-keys.c"

#include "map.c"
#include "list.c"
#include "copy.c"
#include "probes.c"

MODULE_DESCRIPTION("SystemTap probe: shellsnoop");
MODULE_AUTHOR("Martin Hunt <hunt@redhat.com>");

MAP pids, arglist ;

int inst_do_execve (char * filename, char __user *__user *argv, char __user *__user *envp, struct pt_regs * regs)
{
  struct map_node *ptr;

  /* watch shells only */
  /* FIXME: detect more shells, like csh, tcsh, zsh */
  
  if (!strcmp(current->comm,"bash") || !strcmp(current->comm,"sh") || !strcmp(current->comm, "zsh")
      || !strcmp(current->comm, "tcsh") || !strcmp(current->comm, "pdksh"))
    {
      _stp_printf ("%d\t%d\t%d\t%s ", current->uid, current->pid, current->parent->pid, filename);

      _stp_map_key_int64 (pids, current->pid);
      _stp_map_set_int64 (pids, 1);
      
      _stp_list_clear (arglist);
      _stp_copy_argv_from_user (arglist, argv);
      
      foreach (arglist, ptr)
	_stp_printf ("%s ", _stp_get_str(ptr));
      _stp_print("\n");

      _stp_print_flush();
    }
  jprobe_return();
  return 0;
}

struct file * inst_filp_open (const char * filename, int flags, int mode)
{
  _stp_map_key_int64 (pids, current->pid);
  if (_stp_map_get_int64 (pids))
    _stp_printf ("%d\t%d\t%s\tO %s\n", current->pid, current->parent->pid, current->comm, filename);

  _stp_print_flush();
  jprobe_return();
  return 0;
}

asmlinkage ssize_t inst_sys_read (unsigned int fd, char __user * buf, size_t count)
{
  _stp_map_key_int64 (pids, current->pid);
  if (_stp_map_get_int64 (pids)) {
    _stp_printf ("%d\t%d\t%s\tR %d\n", current->pid, current->parent->pid, current->comm, fd);    
    _stp_print_flush();
  }
  jprobe_return();
  return 0;
}

asmlinkage ssize_t inst_sys_write (unsigned int fd, const char __user * buf, size_t count)
{
  _stp_map_key_int64 (pids, current->pid);
  if (_stp_map_get_int64 (pids))
    {
      String str = _stp_string_init (0);
      _stp_string_from_user(str, buf, count);
      _stp_printf ("%d\t%d\t%s\tW %s", current->pid, current->parent->pid, 
		   current->comm, _stp_string_ptr(str));
      _stp_print_flush();
    }
  
  jprobe_return();
  return 0;
}

static struct jprobe stp_probes[] = {
  {
    .kp.addr = (kprobe_opcode_t *)"do_execve",
    .entry = (kprobe_opcode_t *) inst_do_execve
  },
  {
    .kp.addr = (kprobe_opcode_t *)"filp_open",
    .entry = (kprobe_opcode_t *) inst_filp_open
  },
  {
    .kp.addr = (kprobe_opcode_t *)"sys_read",
    .entry = (kprobe_opcode_t *) inst_sys_read
  },
  {
    .kp.addr = (kprobe_opcode_t *)"sys_write",
    .entry = (kprobe_opcode_t *) inst_sys_write
  }, 
};

#define MAX_STP_ROUTINE (sizeof(stp_probes)/sizeof(struct jprobe))


int probe_start(void)
{
	int ret;

	/* now initialize any data or variables */
	pids = _stp_map_new_int64 (10000, INT64);
	arglist = _stp_list_new (10, STRING);

	/* now we are ready to enable the probes */
	ret = _stp_register_jprobes (stp_probes, MAX_STP_ROUTINE);
	
	if (ret < 0) {
	  _stp_map_del (pids);
	  _stp_map_del (arglist);
	  return ret;
	}
	
	_stp_printf("instrumentation is enabled... %s\n", __this_module.name);
	_stp_print_flush();
	return ret;
}


void probe_exit (void)
{
	_stp_unregister_jprobes (stp_probes, MAX_STP_ROUTINE);
	_stp_map_del (pids);
	_stp_map_del (arglist);
	_stp_printf("\nDropped %d packets\n", atomic_read(&_stp_transport_failures));
	_stp_print_flush();
}