/*
 *  jail.c - Ruby jParallel
 *
 *  vim: set nosta noet ts=4 sw=4:
 *
 *  $Id$
 *  
 *  Authors:
 *	* Michael Granger <ged@FaerieMUD.org>
 *	* Mahlon E. Smith <mahlon@martini.nu>
 *  
 *  Copyright (c) 2008, Michael Granger and Mahlon E. Smith
 *  All rights reserved.
 *  
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions are met:
 *  
 *	* Redistributions of source code must retain the above copyright notice,
 *	  this list of conditions and the following disclaimer.
 *  
 *	* Redistributions in binary form must reproduce the above copyright notice,
 *	  this list of conditions and the following disclaimer in the documentation
 *	  and/or other materials provided with the distribution.
 *  
 *	* Neither the name of the author/s, nor the names of the project's
 *	  contributors may be used to endorse or promote products derived from this
 *	  software without specific prior written permission.
 *  
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
 *  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 *  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 *  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 *  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 *  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *  
 */

#include "jail.h"


/* --------------------------------------------------------------
 * Utility functions
 * -------------------------------------------------------------- */

/*
 *  Debug logging function
 */
void
#ifdef HAVE_STDARG_PROTOTYPES
rbjail_debug(const char *fmt, ...)
#else
rbjail_debug( const char *fmt, va_dcl )
#endif
{
	char buf[BUFSIZ], buf2[BUFSIZ];
	va_list args;

	if ( !RTEST(ruby_debug) ) return;

	snprintf( buf, BUFSIZ, "Debug>>> %s", fmt );

	va_init_list( args, fmt );
	vsnprintf( buf2, BUFSIZ, buf, args );
	fputs( buf2, stderr );
	fputs( "\n", stderr );
	fflush( stderr );
	va_end( args );
}


/*
 * Object validity checker. Returns the data pointer.
 */
static struct xprison *
rbjail_check_jail( VALUE self ) {
	debugMsg(( "Checking a BSD::Jail object (%d).", self ));
	Check_Type( self, T_DATA );

	if ( !rb_obj_is_kind_of(self, rbjail_cBSDJail) ) {
		rb_raise( rb_eTypeError, "wrong argument type %s (expected BSD::Jail)",
			rb_obj_classname( self ));
	}

	return DATA_PTR( self );
}


/*
 * Fetch the data pointer and check it for sanity.
 */
static struct xprison *
rbjail_get_jailptr( VALUE self ) {
	struct xprison *ptr = rbjail_check_jail( self );

	debugMsg(( "Fetching a Jail (%p).", ptr ));
	if ( !ptr )
		rb_raise( rb_eRuntimeError, "uninitialized Jail" );

	return ptr;
}


/*
 * Attach to a running jail and chdir to the root.
 */
static int
rbjail_do_attach( int jid )
{
	int attach_status = jail_attach(jid);

	if ( attach_status == -1 )
		rb_sys_fail( "jail_attach" );
	if ( chdir("/") == -1 ) rb_sys_fail( "chdir" );

	return attach_status;
}


/*
 * Attach to a running jail from within a block, forking first
 * and returning the child pid.
 */
static VALUE
rbjail_attach_block( int jid )
{
	int pid;
	int status;

	if ( ! rb_block_given_p() ) return Qnil;

	rb_secure(2);

	fflush(stdout);
	fflush(stderr);

	switch ( pid = fork() ) {
		case 0:
			rb_thread_atfork();
			rbjail_do_attach( jid );
			rb_protect( rb_yield, Qundef, &status );
			rb_exit( status );
			return Qnil;

		case -1:
			rb_sys_fail( "fork(2)" );
			return Qnil;

		default:
			return INT2FIX( pid );
	}
}



/* --------------------------------------------------------------
 * Memory-management functions
 * -------------------------------------------------------------- */

/*
 * Copy memory from the given 'xp' to a ruby managed object.
 */
static VALUE
rbjail_alloc( VALUE class, struct xprison *xp )
{
	struct xprison *rbjail_xp = ALLOC( struct xprison );
	VALUE rbjail = rb_funcall( class, rb_intern("allocate"), 0 );

	// replace the null pointer obj with an xprison ptr.
	//
	memcpy( rbjail_xp, xp, sizeof( struct xprison ) );
	DATA_PTR( rbjail ) = rbjail_xp;

	return rbjail;
}


/*
 * GC Free function
 */
static void
rbjail_gc_free( struct xprison *ptr ) {
	if ( ptr ) {
		xfree( ptr );
	}

	else {
		debugMsg(( "Not freeing an uninitialized jail" ));
	}
}


/*
 * Allocate a new ruby object with a NULL pointer.
 */
static VALUE
rbjail_s_alloc( VALUE class )
{
	return Data_Wrap_Struct( class, NULL, rbjail_gc_free, NULL );
}


/* --------------------------------------------------------------
 * Class methods
 * -------------------------------------------------------------- */

/*
 *  call-seq:
 *     BSD::Jail.new( ip, path, host )   => jail_id
 *
 *  Create a new BSD::Jail object from required +ip+, +path+, and +host+
 *  arguments.  You can optionally pass a +securelevel+ fourth argument.
 *
 *  Returns the +id+ of the newly created jail.
 */
static VALUE
rbjail_jail( int argc, VALUE *argv, VALUE self )
{
	struct jail j;
	struct in_addr in;
	VALUE ip, path, host, sec_level;
	int id;
	int securelevel = -1;

	rb_scan_args( argc, argv, "31", &ip, &path, &host, &sec_level );
	if ( sec_level != Qnil ) securelevel = FIX2INT( sec_level );

	if ( inet_aton( RSTRING_PTR( rb_obj_as_string(ip) ), &in ) == 0 )
		rb_raise( rb_eArgError, "Could not make sense of ip number: %s", ip );

	SafeStringValue(path);
	SafeStringValue(host);

	j.version   = 0;
	j.path	    = RSTRING_PTR( path );
	j.hostname  = RSTRING_PTR( host );
	j.ip_number = ntohl( in.s_addr );
	id = jail(&j);

	if ( id == -1 ) rb_sys_fail( "jail" );
	if ( chdir("/") == -1 ) rb_sys_fail( "chdir" );

	if ( securelevel > -1 && securelevel < 4 ) {
		debugMsg(( "Setting securelevel to: %d", securelevel ));
		if ( sysctlbyname("kern.securelevel", NULL, 0, &securelevel, sizeof(securelevel)) )
			rb_sys_fail( "securelevel" );
	}

	debugMsg(( "New jail created with id: %d\n", id ));
	return INT2FIX( id );
}


/*
 *  call-seq:
 *     BSD::Jail.find_by_jid( jid ) => BSD::Jail
 *
 *  Iterate over the currently instantiated jails, returning a jail
 *  object that matches the given +jid+, or nil if no jail is found.
 *
 */
static VALUE
rbjail_find_by_jid( VALUE self, VALUE jid )
{
	VALUE args[0];

	if ( TYPE(jid) != T_FIXNUM )
		rb_raise( rb_eTypeError, "invalid argument to find_by_jid(): %s",
			RSTRING_PTR( rb_inspect(jid)) );

	args[0] = jid;
	return rbjail_find( 1, args, self );
}


/*
 *  call-seq:
 *     BSD::Jail.search( id )                      => BSD::Jail
 *     BSD::Jail.search( hostname )                => [ BSD::Jail, ... ]
 *     BSD::Jail.search( IPAddr.new('127.0.0.1') ) => [ BSD::Jail, ... ]
 *     BSD::Jail.search( Pathname.new('/tmp') )    => [ BSD::Jail, ... ]
 *     BSD::Jail.search                            => [ BSD::Jail, BSD::Jail, ... ]
 *     BSD::Jail.search { |obj| block }            => obj
 *
 *  Iterate over the currently instantiated jails, returning a jail
 *  object that matches the given argument. 
 *
 *  If the argument is an integer, it is assumed to be a JID.
 *  Otherwise, it is converted to a string, and compared against
 *  IP addresses, hostnames, and paths.  JIDs are unique, so there is
 *  only one object that can match.  Only the matching object (or nil)
 *  is returned in that event.  IPs, hostnames, and paths can be
 *  shared between jails, so searching on those return an array populated
 *  with matched jail objects -- or if there are no valid matches, an empty array.
 *
 *  Without an argument, return an array of all jails as objects.
 *
 */
static VALUE
rbjail_find( int argc, VALUE *argv, VALUE self )
{ 
	struct xprison *sxp, *xp;
	struct in_addr in;
	size_t i, len;

	VALUE arg, rbjail;
	VALUE jails = rb_ary_new();
	int jid = 0, compare = 0;
	char *str = "";

	rb_scan_args( argc, argv, "01", &arg );

	// An argument was passed, so let's figure out what it was
	// and try to compare it to the current jails.
	//
	if ( argc == 1 ) {
		switch ( TYPE(arg) ) {

			// find by JID
			//
			case T_FIXNUM:
				jid = FIX2INT( arg );
				break;

			// find by IP/hostname/path
			//
			case T_OBJECT:
			case T_DATA:
			case T_STRING:
				str = RSTRING_PTR( rb_obj_as_string(arg) );
				compare = 1;
				break;

			default:
				rb_raise( rb_eTypeError, "invalid argument: %s",
						RSTRING_PTR( rb_inspect(arg)) );
		}
	}

	// Get the size of the xprison and allocate memory to it.
	//
	if ( sysctlbyname("security.jail.list", NULL, &len, NULL, 0) == -1 )
		rb_sys_fail("sysctlbyname(): security.jail.list");
	if ( len <= 0 ) {
		rb_sys_fail("sysctlbyname(): unable to determine xprison size");
		return Qnil;
	}

	sxp = xp = malloc( len );
	if ( sxp == NULL ) {
		rb_sys_fail("sysctlbyname(): unable to allocate memory");
		return Qnil;
	}

	// Get and sanity check the current prison list
	//
	if ( sysctlbyname("security.jail.list", xp, &len, NULL, 0) == -1 ) {
		if ( errno == ENOMEM ) free( sxp );
		rb_sys_fail("sysctlbyname(): out of memory");
		return Qnil;
	}
	if ( len < sizeof(*xp) || len % sizeof(*xp) || xp->pr_version != XPRISON_VERSION )
		rb_fatal("Kernel and userland out of sync");

	// No arguments to find() -- yield each successive jail,
	// and return an array of all jail objects.
	//
	if ( argc == 0 ) {
		for ( i = 0; i < len / sizeof(*xp); i++ ) {
			rbjail = rbjail_alloc( self, xp );
			if ( rb_block_given_p() ) rb_yield( rbjail );
			rb_ary_push( jails, rbjail_alloc( self, xp ) );
			xp++;
		}

		free( sxp );
		return jails;
	}

	// Argument passed to find(): walk the jail list, comparing the arg
	// with each current jail.  Return all matches.
	//
	for ( i = 0; i < len / sizeof(*xp); i++ ) {
		in.s_addr = ntohl(xp->pr_ip);
		if (( compare == 0 && xp->pr_id == jid ) ||
				( compare == 1 &&
				  (( strcmp( str, inet_ntoa(in) ) == 0 ) ||
				   ( strcmp( str, xp->pr_host )   == 0 ) ||
				   ( strcmp( str, xp->pr_path )   == 0 ))
				)) {

			debugMsg(( "Located jail: %d", xp->pr_id ));

			// If the user already knows the JID, there can only be
			// one match.  Return it immediately.
			//
			if ( compare == 0 ) {
				free( sxp );
				return rbjail_alloc( self, xp );
			}

			// If searching for anything other than JID, the argument
			// isn't unique.  Put matching jails into an array.
			//
			rb_ary_push( jails, rbjail_alloc( self, xp ) );
			xp++;
		}
		else {

			xp++;
		}
	}

	free( sxp );

	if ( compare == 0 && rb_ary_shift( jails ) == Qnil )
		return Qnil;

	return jails;
}


/*
 *  call-seq:
 *     BSD::Jail.attach( id )       => true
 *     BSD::Jail.attach( hostname ) => true
 *
 *  Attach to a currently running jail instance directly.
 *  Operates under the same rules as BSD::Jail#search -- you may
 *  specify a jail by IP Address, hostname, or jid.
 *
 *  Please note that attaching your process into a jail is a one way
 *  operation that requires root privileges.  You must fork() if
 *  your process needs to continue in the host environment.
 *
 */
static VALUE
rbjail_class_attach( VALUE self, VALUE arg )
{
	int jid = 0;
	VALUE jails;
	VALUE find_args [0];

	switch ( TYPE(arg) ) {

		// The user knows the JID already, attach directly.
		//
		case T_FIXNUM:

			jid = FIX2INT( arg );
			rbjail_do_attach( jid );
			break;

		// Find the JID to attach to.  First match wins.
		//
		case T_OBJECT:
		case T_DATA:
		case T_STRING:

			find_args[0] = arg;
			jails = rbjail_find( 1, find_args, self );
			jid = FIX2INT( rbjail_get_jid( rb_ary_shift(jails) ));
			rbjail_do_attach( jid );
			break;

		default:
			rb_raise( rb_eTypeError, "invalid argument to attach(): %s",
						RSTRING_PTR( rb_inspect(arg)) );
	}

	return Qtrue;
}


/* --------------------------------------------------------------
 * Instance methods
 * -------------------------------------------------------------- */

/*
 *  call-seq:
 *     BSD::Jail <=> BSD::Jail
 *
 *  Interface for Comparable.
 *
 */
static VALUE
rbjail_compare( VALUE self, VALUE other )
{
	debugMsg(("self id: %s, other id: %s",
				RSTRING_PTR(rb_inspect( self )),
				RSTRING_PTR(rb_inspect( other ))));

	return rb_funcall( 
		rbjail_get_jid( self ),
		rb_intern("<=>"),
		1,
		rbjail_get_jid( other )
	);
}


/*
 * Fetch the configured IP address for the jail.
 * Returns an IPAddr object.
 *
 */
static VALUE
rbjail_get_ip( VALUE self )
{
	struct xprison *xp = rbjail_get_jailptr( self );
	struct in_addr in;
	VALUE args [0];

	in.s_addr = ntohl( xp->pr_ip );
	args[0] = rb_str_new2( inet_ntoa(in) );

	return rb_class_new_instance( 1, args, rbjail_cIPAddr );
}


/*
 * Fetch the assigned JID for the jail.
 *
 */
static VALUE
rbjail_get_jid( VALUE self )
{
	struct xprison *xp = rbjail_get_jailptr( self );
	return INT2FIX( xp->pr_id );
}


/*
 * Fetch the configured hostname for the jail as a string.
 *
 */
static VALUE
rbjail_get_host( VALUE self )
{
	struct xprison *xp = rbjail_get_jailptr( self );
	return rb_str_new2( xp->pr_host );
}


/*
 * Fetch the configured path for the jail.
 * Returns a Pathname object.
 *
 */
static VALUE
rbjail_get_path( VALUE self )
{
	struct xprison *xp = rbjail_get_jailptr( self );
	VALUE args[0];

	args[0] = rb_str_new2( xp->pr_path );

	return rb_class_new_instance( 1, args, rbjail_cPathname );
}


/*
 * Return a human readable version of the object.
 */
static VALUE
rbjail_inspect( VALUE self )
{
	char inspect_str[BUFSIZ];

	sprintf( inspect_str, "#<%s:0x%07x jid:%d (%s)>",
		rb_obj_classname( self ),
		NUM2UINT(rb_obj_id( self )) * 2,
		FIX2INT(rbjail_get_jid( self )),
		RSTRING_PTR(rbjail_get_host( self ))
	);

	return rb_str_new2( inspect_str );
}


/*
 *  call-seq:
 *     jail.attach
 *     jail.attach do ... end  => child pid
 *
 *  Attach to the jail.
 *
 *  Please note that attaching your process into a jail is a one way
 *  operation that requires root privileges.  You must fork() if
 *  your process needs to continue in the host environment.
 *
 * 	In the block form, a fork() is performed automatically.
 *
 */
static VALUE
rbjail_instance_attach( VALUE self )
{
	int jid = FIX2INT( rbjail_get_jid( self ) );

	if ( rb_block_given_p() ) {
		return rbjail_attach_block( jid );
	}

	else {
		rbjail_do_attach( jid );
	}

	return Qtrue;
}


/* --------------------------------------------------------------
 * Initalizer
 * -------------------------------------------------------------- */

/*
 *  BSD::Jail
 *
 *  Ruby bindings for the FreeBSD jail(2) subsystem.
 *
 *  Example usage:
 *
 *  	require 'bsd/jail'
 *
 *      # create a new jail
 *      jid = BSD::Jail.create( '127.0.0.1', '/tmp', 'testjail' )
 *
 *      # find existing jail(s)
 *      jail  = BSD::Jail.find_by_jid( jid )
 *      jails = BSD::Jail.search( hostname )
 *      jails = BSD::Jail.search( IPAddr.new('127.0.0.1') )
 *      jails = BSD::Jail.search( Pathname.new('/tmp') )
 *
 *      # attach to jails
 *      BSD::Jail.attach( jid )
 *      jail.attach do
 *     	    ... do something fancy!
 *      end
 *
 *  BSD::Jail includes behaviors from Enumerable and Comparable,
 *  so you can do things such as the following:
 *
 *      # Are these two instantiated jails the same jid?
 *      BSD::Jail.search( '/tmp' ) == BSD::Jail.search( 'testjail' )
 *      
 *      # Alternative interface for finding a jail by jid
 *      BSD::Jail.find { |j| j.jid == 3 }
 *      
 *      # Return all jails as an array
 *      jails = BSD::Jail.find_all
 *      
 *      # What jails have IPs within the 192.168.16.0/24 class C netblock?
 *      nb = IPAddr.new( '192.168.16.0/24' );
 *      jails = BSD::Jail.find? { |j| nb.include?(j.ip) }
 *
 */
void
Init_jail( void )
{
	// namespacing
	//
	rbjail_mBSD	= rb_define_module( "BSD" );
	rbjail_cBSDJail = rb_define_class_under( rbjail_mBSD, "Jail", rb_cObject );

	// struct wrapping function
	rb_define_alloc_func( rbjail_cBSDJail, rbjail_s_alloc );

	// Make the 'new' method private.
	rb_funcall( rbjail_cBSDJail, rb_intern("private_class_method"), 1, ID2SYM(rb_intern("new")) );

	// class methods
	//
	rb_define_singleton_method( rbjail_cBSDJail, "attach", rbjail_class_attach, 1 );
	rb_define_singleton_method( rbjail_cBSDJail, "create", rbjail_jail, -1 );
	rb_define_singleton_method( rbjail_cBSDJail, "search", rbjail_find, -1 );
	rb_define_singleton_method( rbjail_cBSDJail, "find_by_jid", rbjail_find_by_jid, 1 );

	// instance methods
	//
	rb_define_method( rbjail_cBSDJail, "attach", rbjail_instance_attach, 0 );
	rb_define_method( rbjail_cBSDJail, "host", rbjail_get_host, 0 );
	rb_define_alias(  rbjail_cBSDJail, "hostname", "host" );
	rb_define_method( rbjail_cBSDJail, "inspect", rbjail_inspect, 0 );
	rb_define_method( rbjail_cBSDJail, "ip",   rbjail_get_ip,   0 );
	rb_define_method( rbjail_cBSDJail, "jid",  rbjail_get_jid,  0 );
	rb_define_method( rbjail_cBSDJail, "path", rbjail_get_path, 0 );

	// Additional (base) modules for accessor wrapping
	//
	rb_require( "ipaddr" );
	rb_require( "pathname" );
	rbjail_cIPAddr   = rb_const_get( rb_cObject, rb_intern("IPAddr") );
	rbjail_cPathname = rb_const_get( rb_cObject, rb_intern("Pathname") );

	// Enumerable!
	//
	rb_define_singleton_method( rbjail_cBSDJail, "each", rbjail_find, -1 );
	rb_extend_object( rbjail_cBSDJail, rb_const_get( rb_cObject, rb_intern("Enumerable") ));

	// Comparable!
	//
	rb_define_method( rbjail_cBSDJail, "<=>", rbjail_compare, 1 );
	rb_include_module( rbjail_cBSDJail, rb_const_get( rb_cObject, rb_intern("Comparable") ));
}