#import <Foundation/Foundation.h>
#import <PortMapper/PortMapper.h>
#include <unistd.h>

void printUsageAndExit()
{
    printf("Usage: portmap [-p | -a port | -d port] [-t | -u]\n");
    printf("\t-p\tPrints the external IP Address\n");
    printf("\t-a port\tAdds a new port mapping with the given port\n");
    printf("\t-d port\tDeletes a port mapping with the given port\n");
    printf("\t-t\tSpecify TCP\n");
    printf("\t-u\tSpecify UDP\n");
    exit(1);
}

int readPortAndExitOnError(char *str)
{
    int port = strtol(str, NULL, 10);
    if(port == 0 && errno == EINVAL)
    {
        printf("Invalid port number: %s\n", str);
        printUsageAndExit();
    }
    
    return port;
}

int main (int argc, char * argv[]) {
    NSAutoreleasePool * pool = [[NSAutoreleasePool alloc] init];
    
    BOOL displayRemoteIP = NO;
    BOOL addPortMapping = NO;
    BOOL deletePortMapping = NO;
    int port = -1;
    int protocal = kPMTCP_Protocol;
    int ch;
    
    // parse the arguments
    while((ch = getopt(argc, argv, "pa:d:tu")) != -1)
    {
        switch(ch)
        {
            case 'p':
                displayRemoteIP = YES;
                break;
            case 'a':
                addPortMapping = YES;
                port = readPortAndExitOnError(optarg);
                break;
            case 'd':
                deletePortMapping = YES;
                port = readPortAndExitOnError(optarg);
                break;
            case 't':
                protocal = kPMTCP_Protocol;
                break;
            case 'u':
                protocal = kPMUDP_Protocol;
                break;
            default:
                printUsageAndExit();
                break;
        }
    }
    
    if(addPortMapping && deletePortMapping)
    {
        printf("Can't add and delete a port mapping at the same time");
        printUsageAndExit();
    }
    else if(!displayRemoteIP && !addPortMapping && !deletePortMapping)
    {
        printUsageAndExit();
    }
    
    int result = PMInitialize();
    if(result == kPMNoRouterFoundError)
    {
        printf("Couldn't find a suitable router.\n");
        return -2;
    }
    else if(result == kPMInitializationError)
    {
        printf("Found a router, but was unable to communicate properly with it.\n");
        return -3;
    }
    
    // allow our port mappings to exist after we exist
    PMSetRemoveAtExit(NO);
    
    if(displayRemoteIP)
    {
        NSString* str = (NSString*)PMGetPublicIPAddress();
        
        printf("External IP: %s\n", [str cString]);
        
        [str release];
    }
    
    if(addPortMapping)
    {
        int error = PMAddPortMapping(port, port, protocal);
        if(error == kPMSuccessful)
        {
            printf("Successfully mapped port %d to this machine.\n", port);
        }
        else if(error == kPMPortMappingInUseError)
        {
            printf("Failed to map the port %d to this machine because it is already mapped to a different IP address\n", port);
        }
        else
        {
            printf("Failed to map the port %d to this machine due to an unexpected error.\n", port);
        }
    }
    
    if(deletePortMapping)
    {
        int error = PMRemovePortMapping(port, port, protocal);
        if(error == kPMSuccessful)
        {
            printf("Successfully unmapped port %d.\n", port);
        }
        else if(error == kPMNoSuchPortMappingError)
        {
            printf("Unable to unmap the port %d because the port mapping doesn't exist.\n", port);
        }
        else
        {
            printf("Failed to unmap the port %d do to an unexpected error (%d).\n", port, error);
        }
    }
    
    
    [pool release];
    return 0;
}