/*
 *
 * Copyright (c) 2011 Apple Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <net/if.h>
#include <System/net/pfvar.h>
#include <string.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <AssertMacros.h>
#include "P2PPacketFilter.h"

#define AIRDROP_ANCHOR_PATH	"com.apple/200.AirDrop"
#define MDNS_ANCHOR_NAME	"Bonjour"
#define MDNS_ANCHOR_PATH	AIRDROP_ANCHOR_PATH "/" MDNS_ANCHOR_NAME

#define PF_DEV_PATH "/dev/pf"
#define BONJOUR_PORT 5353

static int openPFDevice( int * outFD )
{
    int err;
    int fd = open( PF_DEV_PATH, O_RDWR );
    
    if( fd >= 0 )
    {
        err = 0;
        *outFD = fd;
    }
    else
    {
        err = errno;
    }
    
    return err;
}

static int getTicket( int devFD, u_int32_t * outTicket, char * anchorPath )
{
    struct pfioc_trans_e trans_e;
    
    trans_e.rs_num = PF_RULESET_FILTER;
    strlcpy( trans_e.anchor, anchorPath, sizeof( trans_e.anchor ) );
    
    struct pfioc_trans trans;
    
    trans.size = 1;
    trans.esize = sizeof( trans_e );
    trans.array = &trans_e;
    
    int result, ioctlError;
    
    ioctlError = ioctl( devFD, DIOCXBEGIN, &trans );
    if( ioctlError )
    {
        result = errno;
    }
    else
    {
        result = 0;
        *outTicket = trans_e.ticket;
    }
    
    return result;
}

static int commitChange( int devFD, u_int32_t ticket, char * anchorPath )
{
    struct pfioc_trans_e trans_e;
    
    trans_e.rs_num = PF_RULESET_FILTER;
    strlcpy( trans_e.anchor, anchorPath, sizeof( trans_e.anchor ) );
    trans_e.ticket = ticket;
    
    struct pfioc_trans trans;
    
    trans.size = 1;
    trans.esize = sizeof( trans_e );
    trans.array = &trans_e;
    
    int result, ioctlError;
    
    ioctlError = ioctl( devFD, DIOCXCOMMIT, &trans );
    if( ioctlError )
        result = errno;
    else
        result = 0;
    
    return result;
}

static int getPoolTicket( int devFD, u_int32_t * outPoolTicket )
{
    struct pfioc_pooladdr pp;
    
    int result, ioctlError;
    
    ioctlError = ioctl( devFD, DIOCBEGINADDRS, &pp );
    if( ioctlError )
    {
        result = errno;
    }
    else
    {
        result = 0;
        *outPoolTicket = pp.ticket;
    }
    
    return result;
}

static int addRule( int devFD, struct pfioc_rule * pr )
{
    int result, ioctlResult;
    
    ioctlResult = ioctl( devFD, DIOCADDRULE, pr );
    if( ioctlResult )
        result = errno;
    else
        result = 0;
    
    return result;
}

static void initRuleHeader( struct pfioc_rule * pr,
                             u_int32_t ticket,
                             u_int32_t poolTicket,
							 char * anchorPath )
{
    pr->action = PF_CHANGE_NONE;
    pr->ticket = ticket;
    pr->pool_ticket = poolTicket;
    strlcpy( pr->anchor, anchorPath, sizeof( pr->anchor ) );
}

// allow inbound traffice on the Bonjour port (5353) 
static void initBonjourRule( struct pfioc_rule * pr,
                              const char * interfaceName,
                              u_int32_t ticket,
                              u_int32_t poolTicket, 
							  char * anchorPath )
{
    memset( pr, 0, sizeof( *pr ) );
    
    // Header
    initRuleHeader( pr, ticket, poolTicket, anchorPath );
    
    // Rule
    pr->rule.dst.xport.range.port[0] = htons( BONJOUR_PORT );
    pr->rule.dst.xport.range.op = PF_OP_EQ;
    
    strlcpy( pr->rule.ifname, interfaceName, sizeof( pr->rule.ifname ) );
    
    pr->rule.action = PF_PASS;
    pr->rule.direction = PF_IN;
    pr->rule.keep_state = 1;
    pr->rule.af = AF_INET6;
    pr->rule.proto = IPPROTO_UDP;
    pr->rule.extfilter = PF_EXTFILTER_APD;
}

// allow outbound TCP connections and return traffic for those connections
static void initOutboundTCPRule( struct pfioc_rule * pr,
                             const char * interfaceName,
                             u_int32_t ticket,
                             u_int32_t poolTicket, 
							 char * anchorPath )
{
    memset( pr, 0, sizeof( *pr ) );
    
    // Header
    initRuleHeader( pr, ticket, poolTicket, anchorPath );
    
    // Rule
    strlcpy( pr->rule.ifname, interfaceName, sizeof( pr->rule.ifname ) );
    
    pr->rule.action = PF_PASS;
    pr->rule.direction = PF_OUT;
    pr->rule.keep_state = 1;
    pr->rule.proto = IPPROTO_TCP;
}

// allow inbound traffic on the specified port and protocol
static void initPortRule( struct pfioc_rule * pr,
                              const char * interfaceName,
                              u_int32_t ticket,
                              u_int32_t poolTicket,
							  char * anchorPath,
							  u_int16_t port,
							  u_int16_t protocol )
{
    memset( pr, 0, sizeof( *pr ) );
    
    // Header
    initRuleHeader( pr, ticket, poolTicket, anchorPath );
    
    // Rule
	// mDNSResponder passes the port in Network Byte Order, so htons(port) is not required 
    pr->rule.dst.xport.range.port[0] = port;
    pr->rule.dst.xport.range.op = PF_OP_EQ;
    
    strlcpy( pr->rule.ifname, interfaceName, sizeof( pr->rule.ifname ) );
    
    pr->rule.action = PF_PASS;
    pr->rule.direction = PF_IN;
    pr->rule.keep_state = 1;
    pr->rule.af = AF_INET6;
    pr->rule.proto = protocol;
    pr->rule.extfilter = PF_EXTFILTER_APD;
}

// allow inbound traffice on the Bonjour port (5353) and the specified port and protocol
int P2PPacketFilterAddBonjourRuleSet(const char * interfaceName, u_int16_t port, u_int16_t protocol )
{
    int result;
    u_int32_t ticket, poolTicket;
    int devFD = -1;
	char * anchorPath = MDNS_ANCHOR_PATH;
	
    result = openPFDevice( &devFD );
    require( result == 0, exit );
    
    result = getTicket( devFD, &ticket, anchorPath );
    require( result == 0, exit );
    
    result = getPoolTicket( devFD, &poolTicket );
    require( result == 0, exit );
    
    struct pfioc_rule pr;
    
    // allow inbound Bonjour traffice to port 5353 
    initBonjourRule( &pr, interfaceName, ticket, poolTicket, anchorPath);
    
    result = addRule( devFD, &pr );
    require( result == 0, exit );
    
    // open inbound port for service
    initPortRule( &pr, interfaceName, ticket, poolTicket, anchorPath, port, protocol );
    
    result = addRule( devFD, &pr );
    require( result == 0, exit );
    
    // allow outbound TCP connections and return traffic for those connections
    initOutboundTCPRule( &pr, interfaceName, ticket, poolTicket, anchorPath);
    
    result = addRule( devFD, &pr );
    require( result == 0, exit );
    
    result = commitChange( devFD, ticket, anchorPath );
    require( result == 0, exit );
    
exit:
    
    if( devFD >= 0 )
        close( devFD );
    
    return result;
}

int P2PPacketFilterClearBonjourRules()
{
	int result;
	int pfDev = -1;
	u_int32_t ticket;
	char * anchorPath = MDNS_ANCHOR_PATH;
	
	result = openPFDevice( &pfDev );
	require( result == 0, exit );
	
	result = getTicket( pfDev, &ticket, anchorPath );
	require( result == 0, exit );
	
	result = commitChange( pfDev, ticket, anchorPath );
	
exit:
	
	if( pfDev >= 0 )
		close( pfDev );

	return result;
}