-- Enhanced Mirai Detection Dissector for Wireshark
-- Optimized for detecting Mirai traffic patterns

Mirai = Proto("MIRAI-CNC", "Mirai CNC Protocol")

-- Known Mirai characteristics
local MIRAI_PATTERNS = {
    TELNET_USER = "^root$",
    TELNET_PASS = "^(admin|root|123456|password)$",
    HTTP_HEADER = "^Mozilla/5.0 %(compatible; MSIE 10.0; Windows NT 6.1; Trident/6.0%)$",
    SCANNER_PATTERN = "^.+:.+$"
}

-- Known Mirai attack types
local ATTACK_TYPES = {
    [0] = "UDP Flood",
    [1] = "VSE Flood",
    [2] = "DNS Water Torture",
    [3] = "SYN Flood",
    [4] = "ACK Flood",
    [5] = "STOMP Flood",
    [6] = "UDP Generic Flood",
    [7] = "HTTP Flood",
    [8] = "TCP Generic Flood"
}

-- Field definitions optimized for detection
local f = Mirai.fields
f.magic = ProtoField.uint16("mirai.magic", "Magic Header", base.HEX)
f.length = ProtoField.uint16("mirai.length", "Length", base.DEC)
f.attack_duration = ProtoField.uint32("mirai.attack_duration", "Attack Duration (seconds)", base.DEC)
f.attack_type = ProtoField.uint8("mirai.attack_type", "Attack Type", base.DEC, ATTACK_TYPES)
f.num_targets = ProtoField.uint8("mirai.num_targets", "Number of Targets", base.DEC)
f.target_ip = ProtoField.ipv4("mirai.target_ip", "Target IP")
f.target_port = ProtoField.uint16("mirai.target_port", "Target Port", base.DEC)
f.target_subnet = ProtoField.uint8("mirai.target_subnet", "Target Subnet", base.DEC)
f.payload = ProtoField.bytes("mirai.payload", "Payload")
f.telnet_data = ProtoField.string("mirai.telnet_data", "Telnet Data")
f.scanner_data = ProtoField.string("mirai.scanner_data", "Scanner Data")

-- Enhanced expert info for detection
local expert = Proto.expert
f.expert_attack = ProtoExpert.new("mirai.attack.expert", "Mirai Attack Detected",
    expert.group.SECURITY, expert.severity.WARN)
f.expert_scan = ProtoExpert.new("mirai.scan.expert", "Mirai Scanner Activity",
    expert.group.SECURITY, expert.severity.NOTE)
f.expert_cnc = ProtoExpert.new("mirai.cnc.expert", "Potential Mirai CNC Communication",
    expert.group.SECURITY, expert.severity.WARN)
f.expert_bot = ProtoExpert.new("mirai.bot.expert", "Potential Mirai Bot Activity",
    expert.group.SECURITY, expert.severity.WARN)
f.expert_bruteforce = ProtoExpert.new("mirai.bruteforce.expert", "Possible Mirai Bruteforce Attempt",
    expert.group.SECURITY, expert.severity.WARN)

Mirai.experts = {f.expert_attack, f.expert_scan, f.expert_cnc, f.expert_bot, f.expert_bruteforce}

-- Detection helper functions
local function detect_telnet_pattern(payload)
    if payload then
        for line in payload:gmatch("[^\r\n]+") do
            if line:match(MIRAI_PATTERNS.TELNET_USER) or line:match(MIRAI_PATTERNS.TELNET_PASS) then
                return true
            end
        end
    end
    return false
end

local function detect_scanner_pattern(payload)
    if payload and payload:match(MIRAI_PATTERNS.SCANNER_PATTERN) then
        return true
    end
    return false
end

local function detect_http_pattern(payload)
    if payload and payload:match(MIRAI_PATTERNS.HTTP_HEADER) then
        return true
    end
    return false
end

-- Main dissector function
function Mirai.dissector(buffer, pinfo, tree)
    if buffer:len() < 4 then return false end
    
    local is_mirai = false
    local offset = 0
    local msg_length = buffer(offset, 2):uint()
    
    pinfo.cols.protocol = "Mirai CNC"
    local subtree = tree:add(Mirai, buffer(), "Mirai Protocol Analysis")
    
    -- Check for Telnet traffic (port 23)
    if pinfo.dst_port == 23 or pinfo.src_port == 23 then
        local payload = buffer:string()
        if detect_telnet_pattern(payload) then
            subtree:add_proto_expert_info(f.expert_bruteforce, "Possible Mirai Telnet Bruteforce")
            subtree:add(f.telnet_data, buffer(offset):string())
            is_mirai = true
        end
    end
    
    -- Check for scanner activity
    if pinfo.dst_port >= 1024 and pinfo.dst_port <= 65535 then
        local payload = buffer:string()
        if detect_scanner_pattern(payload) then
            subtree:add_proto_expert_info(f.expert_scan, "Mirai Scanner Activity Detected")
            subtree:add(f.scanner_data, buffer(offset):string())
            is_mirai = true
        end
    end
    
    -- Check for CNC communication
    if buffer:len() >= msg_length and msg_length >= 8 then
        subtree:add(f.length, msg_length)
        offset = offset + 2
        
        -- Attack command analysis
        local attack_tree = subtree:add(Mirai, buffer(offset), "Attack Analysis")
        local attack_duration = buffer(offset, 4):uint()
        attack_tree:add(f.attack_duration, attack_duration)
        offset = offset + 4
        
        local attack_type = buffer(offset, 1):uint()
        if ATTACK_TYPES[attack_type] then
            attack_tree:add(f.attack_type, attack_type)
            attack_tree:add_proto_expert_info(f.expert_attack, 
                string.format("Mirai %s Attack Detected", ATTACK_TYPES[attack_type]))
            is_mirai = true
        end
        offset = offset + 1
        
        -- Target analysis
        local num_targets = buffer(offset, 1):uint()
        if num_targets > 0 and num_targets < 255 then
            subtree:add(f.num_targets, num_targets)
            offset = offset + 1
            
            local targets_tree = subtree:add(Mirai, buffer(offset), "Target Analysis")
            for i = 1, num_targets do
                local target_tree = targets_tree:add(Mirai, buffer(offset, 7), 
                    string.format("Target %d", i))
                target_tree:add(f.target_ip, buffer(offset, 4))
                target_tree:add(f.target_port, buffer(offset + 4, 2))
                target_tree:add(f.target_subnet, buffer(offset + 6, 1))
                offset = offset + 7
            end
            is_mirai = true
        end
        
        -- Payload analysis
        if offset < buffer:len() then
            local payload = buffer(offset):string()
            subtree:add(f.payload, buffer(offset))
            
            if detect_http_pattern(payload) then
                subtree:add_proto_expert_info(f.expert_bot, "Mirai Bot HTTP Traffic")
                is_mirai = true
            end
        end
    end
    
    -- Set protocol column if Mirai activity detected
    if is_mirai then
        pinfo.cols.protocol = "MIRAI"
        pinfo.cols.info = "Mirai Botnet Traffic"
        return true
    end
    
    return false
end

-- Register for all relevant ports
local tcp_table = DissectorTable.get("tcp.port")
local udp_table = DissectorTable.get("udp.port")

-- Known Mirai ports
local ports = {
    23,     -- Telnet
    48101,  -- CNC
    53,     -- DNS
    80,     -- HTTP
    443,    -- HTTPS
    37810,  -- Additional CNC
    5555,   -- Common target
    7547,   -- Common target
    35105,  -- Scanner
    61001   -- Scanner
}

-- Register for both TCP and UDP
for _, port in ipairs(ports) do
    tcp_table:add(port, Mirai)
    udp_table:add(port, Mirai)
end
