@load base/protocols/smb

module Cobaltstrike;


export {
    # Number of times to match pattern before logging
    const beacon_iterations: count = 4 &redef;
    # Max size of read/write to track
    const hb_len_limit = 25; 
}

export {
	redef enum Notice::Type +=
	{
        C1::Beacon_Activity
    };
}


# start with buffer overflow response
event smb2_message(c: connection, hdr: SMB2::Header, is_orig: bool) {
    if (c$smb_state$current_cmd$command == "IOCTL" && c$smb_state$current_cmd?$status) {
        if (c$smb_state$current_cmd$status == "BUFFER_OVERFLOW") {
            SumStats::observe("cs_smb",
                      SumStats::Key($host=c$id$orig_h,
                                     $str=cat_sep($sep="|",
                                     $def="NULL",
                                     c$id$orig_p,
                                     c$id$resp_h,
                                     c$id$resp_p)),
                      SumStats::Observation($num=0,
                        $str=cat_sep($sep="|", $def="NULL", "buffer_overflow")));
            }
        }
    }



event smb2_read_request(c: connection, hdr: SMB2::Header, file_id: SMB2::GUID, offset: count, length: count) {
    # if tagged, mark size 
        SumStats::observe("cs_smb",
                  SumStats::Key($host=c$id$orig_h,
                                 $str=cat_sep($sep="|",
                                 $def="NULL",
                                 c$id$orig_p,
                                 c$id$resp_h,
                                 c$id$resp_p)),
                  SumStats::Observation($num=length, 
                    $str=cat_sep($sep="|", $def="NULL", "read", length)));
    }

event smb2_write_request(c: connection, hdr: SMB2::Header, file_id: SMB2::GUID, offset: count, length: count) {

    if (length < hb_len_limit) {
        SumStats::observe("cs_smb",
                      SumStats::Key($host=c$id$orig_h,
                                     $str=cat_sep($sep="|",
                                     $def="NULL",
                                     c$id$orig_p,
                                     c$id$resp_h,
                                     c$id$resp_p)),
                      SumStats::Observation($num=length,
                        $str=cat_sep($sep="|", $def="NULL", "write", length)));
                      }
}

function thresh_crossed(key:SumStats::Key, result:SumStats::Result)
{

    local r = result["cs_smb"];
    #unpack our key 0: orig_p, 1: resp_h, 2: resp_p
    local keys = split_string($str = key$str, $re=/\|/);
  
    # note data_len is a str now because of pack/unpack
    #local s = fmt("Possible SMB Beacon traffic: observered %.0f writes of data length %s", r$sum, keys[2]);
    local s = fmt("Potential Cobalt Strike SMB Beacon Activity.");
    print fmt(s);
    local cid: conn_id = [$orig_h = key$host,
                         $orig_p = to_port(keys[0]),
                         $resp_h = to_addr(keys[1]),
                         $resp_p = to_port(keys[2])];
    
    # TODO: error check this
    local conn = lookup_connection(cid);
    # Raise Notice
    NOTICE([$note=C1::Beacon_Activity,
            $msg=s,
            #$src=key$host,
            #$dst=to_addr(keys[0]),
            #$identifier=key$str,
            $conn = conn,
            $suppress_for=30min]);

    SumStats::next_epoch("CS_SMB_Beacon");
}

function check_beacon(key: SumStats::Key, result: SumStats::Result): double {   
    local cmds = vector("read" ,"write", "buffer_overflow");
    local mark: int = -1;
    local len: count = 0;
    local r_vec = SumStats::get_last(result["cs_smb"]);
    for (i in r_vec)
    {
        local keys = split_string($str = r_vec[i]$str,  $re=/\|/);
        if (mark == -1)
        {
            for (j in cmds)
            {
                if (keys[0] == cmds[j])
                {
                    # keep track of where we are in the sequence
                    mark = j;
                    # get out of this loop
                    break;
                }
            }
        }
        # mark is set if first cmd write or read
        if (len == 0 && r_vec[i]$num != 0)
        {
            len = r_vec[i]$num;
        }
        # len is set
        if (keys[0] == cmds[mark])
        {
            if(cmds[mark] == "read" || cmds[mark] == "write")
            {   
                if(len != r_vec[i]$num)
                {
                    return 1.0;
                }
            }
            else
            {
                if (cmds[mark] != "buffer_overflow")
                {
                    #print fmt("cmd: %s\t\tvec: %s",cmds[mark], keys[0]);
                    return 1.0;
                }
            }
        }
        else
        {
            return 1.0;
        }
        
        # move mark up
        mark = (mark + 1) % 3;

    }

    # winner, fit the pattern all the same length
    return 3.0;
}

event zeek_init()
{

    local r1 = SumStats::Reducer($stream="cs_smb", 
                                 $apply=set(SumStats::LAST),
                                 $num_last_elements = beacon_iterations * 3
                                 );

    
    SumStats::create([$name = "CS_SMB_Beacon",
                      $epoch = 0secs, #manual epochs
                      $reducers = set(r1),
                      $threshold = 2.0,
                      $threshold_val = check_beacon,
                      $threshold_crossed = thresh_crossed
                       ]);
}