// DefenderRuleParser
// Author: Andrea Cristaldi 2025 - https://github.com/andreacristaldi/DefenderRuleParser
// This project is licensed under the Apache 2.0 License.
/*
 * Summary: SIGTREE_EXT parser — extended tree structure variant.
 * Origin: dump-driven; additional flags/attr space seen in dumps guided the implementation.
 * Role: Mirrors the base SIGTREE approach with extra attributes preserved.
 */
using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using DefenderRuleParser2.Models;
using DefenderRuleParser2.Parsers.Wildcards; 

namespace DefenderRuleParser2.Parsers
{
    
    public sealed class SigTreeExtParser : ISignatureParser
    {
        
        private const int MIN_PREVIEW_ASCII_RUN = 6;

        public void Parse(BinaryReader reader, int size, uint threatId)
        {
            long baseAbs = reader.BaseStream.Position;
            long endAbs = baseAbs + size;

            if (size < 4)
            {
                Logger.Warn(string.Format("[!] SIGTREE_EXT: size too small ({0}) @0x{1:X}", size, baseAbs));
                reader.BaseStream.Seek(endAbs, SeekOrigin.Begin);
                return;
            }

            var buf = reader.ReadBytes(size);
            if (buf == null || buf.Length != size)
            {
                Logger.Warn(string.Format("[!] SIGTREE_EXT: short read ({0}/{1}) @0x{2:X}",
                    (buf == null ? 0 : buf.Length), size, baseAbs));
                reader.BaseStream.Seek(endAbs, SeekOrigin.Begin);
                return;
            }

            int p = 0;
            if (p + 4 > buf.Length)
            {
                Logger.Warn(string.Format("[!] SIGTREE_EXT: header truncated @0x{0:X}", baseAbs));
                reader.BaseStream.Seek(endAbs, SeekOrigin.Begin);
                return;
            }

            
            ushort nodeCount = ReadU16(buf, ref p);
            ushort flags = ReadU16(buf, ref p);

            Logger.Info(string.Format("[SIGTREE_EXT] Threat={0} @0x{1:X}, Size={2}, Nodes={3}, Flags=0x{4:X}",
                                      threatId, baseAbs, size, nodeCount, flags));

            var nodes = new List<string>(Math.Max(8, (int)nodeCount));
            int parsedNodes = 0;

            
            int firstMarker = FindNextMarker(buf, p);
            if (firstMarker > p)
            {
                int prologLen = firstMarker - p;
                bool nonZero = HasAnyNonZero(buf, p, prologLen);

                Logger.Info(string.Format("  - [PROLOG] @+0x{0:X} len={1}", p, prologLen));
                LogPayloadPreviews(buf, p, prologLen, "    ");

                if (nonZero && parsedNodes < nodeCount)
                {
                    nodes.Add(string.Format("Node#{0}: PROLOG len={1}", parsedNodes, prologLen));
                    parsedNodes++;
                }

                p = firstMarker;
            }

            
            while (parsedNodes < nodeCount && p + 2 <= buf.Length)
            {
                
                if (!IsMarkerAt(buf, p))
                {
                    int next = FindNextMarker(buf, p + 1);
                    if (next < 0) break;
                    p = next;
                }

                int nodeStart = p;                
                long nodeStartAbs = baseAbs + nodeStart;

                p += 2; 

                if (p + 9 > buf.Length)
                {
                    Logger.Warn(string.Format("[SIGTREE_EXT] Node#{0} header truncated @+0x{1:X} (abs=0x{2:X})",
                                              parsedNodes, nodeStart, nodeStartAbs));
                    break;
                }

                
                byte nType = buf[p++];
                byte nFlags = buf[p++];
                ushort child1 = ReadU16(buf, ref p);
                ushort child2 = ReadU16(buf, ref p);
                ushort weight = ReadU16(buf, ref p);
                byte u8 = buf[p++];

                
                int payloadStart = p;
                int payloadEnd;
                if (parsedNodes == nodeCount - 1)
                {
                    
                    payloadEnd = buf.Length;
                }
                else
                {
                    int nextMarker = FindNextMarker(buf, p);
                    payloadEnd = (nextMarker >= 0) ? nextMarker : buf.Length;
                }
                if (payloadEnd < payloadStart) payloadEnd = payloadStart;
                int payloadLen = payloadEnd - payloadStart;

                Logger.Info(string.Format("  - [#{0}] @0x{1:X} len={2}  Type=0x{3:X2} ({4}) Flags=0x{5:X2} Child1={6} Child2={7} Weight={8} U8=0x{9:X2}",
                                          parsedNodes, nodeStartAbs, (payloadEnd - nodeStart),
                                          nType, DescribeLogic(nType), nFlags, child1, child2, weight, u8));

                Logger.Info(string.Format("      Payload: rel=+0x{0:X}, abs=0x{1:X}, len={2}",
                                          payloadStart, baseAbs + payloadStart, payloadLen));

                if (payloadLen > 0)
                {
                    LogPayloadPreviews(buf, payloadStart, payloadLen, "      ");
                }
                else
                {
                    Logger.Info("      (empty payload)");
                }

                nodes.Add(string.Format("Node#{0}: type={1} w={2} c1={3} c2={4} payloadLen={5}",
                                        parsedNodes, DescribeLogic(nType), weight, child1, child2, payloadLen));

                p = payloadEnd;
                parsedNodes++;
            }

            
            if (p < buf.Length)
            {
                int tailStart = p;
                int tailLen = buf.Length - tailStart;
                Logger.Info(string.Format("  - [TAIL] Unparsed bytes after nodes: rel=+0x{0:X}, len={1}", tailStart, tailLen));
                LogPayloadPreviews(buf, tailStart, tailLen, "      ");
            }

            
            if (parsedNodes != nodeCount)
                Logger.Warn(string.Format("[SIGTREE_EXT] ParsedNodes={0}, Expected={1}", parsedNodes, nodeCount));
            else
                Logger.Info(string.Format("[SIGTREE_EXT] ParsedNodes={0}, Expected={1} (OK)", parsedNodes, nodeCount));

            
            try
            {
                if (ThreatDatabase.TryGetThreat(threatId, out var threat))
                {
                    threat.Signatures.Add(new SignatureEntry
                    {
                        Type = "SIGNATURE_TYPE_SIGTREE_EXT",
                        Offset = baseAbs,
                        Parsed = (parsedNodes == nodeCount),
                        Pattern = nodes,
                        ConditionType = "TREE_STRUCTURE",
                        ConditionValue = nodeCount,
                        Logic = new SignatureLogic { Threshold = nodeCount }
                    });
                }
            }
            catch (Exception ex)
            {
                Logger.Warn(string.Format("[SIGTREE_EXT] ThreatDatabase update failed: {0}", ex.Message));
            }

            
            reader.BaseStream.Seek(endAbs, SeekOrigin.Begin);
        }

        

        private static ushort ReadU16(byte[] b, ref int p)
        {
            if (p + 2 > b.Length) return 0;
            ushort v = (ushort)(b[p] | (b[p + 1] << 8));
            p += 2;
            return v;
        }

        private static bool IsMarkerAt(byte[] b, int ofs)
        {
            return (ofs + 1 < b.Length) && b[ofs] == 0x30 && b[ofs + 1] == 0x00;
        }

        private static int FindNextMarker(byte[] b, int start)
        {
            for (int i = Math.Max(0, start); i + 1 < b.Length; i++)
                if (b[i] == 0x30 && b[i + 1] == 0x00)
                    return i;
            return -1;
        }

        private static bool HasAnyNonZero(byte[] b, int start, int len)
        {
            int end = Math.Min(b.Length, start + len);
            for (int i = start; i < end; i++)
                if (b[i] != 0x00) return true;
            return false;
        }

        private static string Hex(byte[] b, int start, int len)
        {
            int n = Math.Min(len, b.Length - start);
            if (n < 0) n = 0;
            var sb = new StringBuilder(n * 3);
            for (int i = 0; i < n; i++)
            {
                sb.Append(b[start + i].ToString("X2"));
                if (i + 1 < n) sb.Append(' ');
            }
            return sb.ToString();
        }

        private static bool IsPrintable(byte x) { return x >= 0x20 && x <= 0x7E; }

        
        private static void FindLongestAsciiRun(byte[] b, int start, int len, out int relStart, out int runLen)
        {
            relStart = -1; runLen = 0;
            int end = Math.Min(b.Length, start + len);
            int i = start;
            while (i < end)
            {
                
                while (i < end && !IsPrintable(b[i])) i++;
                if (i >= end) break;

                int j = i;
                while (j < end && IsPrintable(b[j])) j++;

                int curLen = j - i;
                if (curLen > runLen)
                {
                    runLen = curLen;
                    relStart = i - start;
                }

                i = j;
            }

            if (runLen < MIN_PREVIEW_ASCII_RUN)
            {
                relStart = -1;
                runLen = 0;
            }
        }

        
        private static int ConsumeWildcardBlock(byte[] b, int start, int len, int iAbs)
        {
            int endAbs = Math.Min(b.Length, start + len);
            int i = iAbs;
            if (i >= endAbs || b[i] != 0x90) return iAbs;

            if (i + 1 >= endAbs) return iAbs;
            byte op = b[i + 1];

            if (op == 0x00) return iAbs; 
            if (op == 0x90) return (i + 2); 

            if (op >= 0x32)
            {
                
                return iAbs;
            }

            
            if ((op >= 0x01 && op <= 0x03) || op == 0x17)
            {
                if (i + 2 < endAbs) return (i + 3);
                return iAbs;
            }

            
            if (op == 0x04 || op == 0x05 || op == 0x19 || op == 0x1A)
            {
                if (i + 3 < endAbs) return (i + 4);
                return iAbs;
            }

            
            if (op == 0x09)
            {
                if (i + 3 < endAbs) return (i + 4);
                return iAbs;
            }

            
            if (op == 0x1B)
            {
                if (i + 1 < endAbs) return (i + 2);
                return iAbs;
            }

            
            if (i + 3 < endAbs)
            {
                int lo = b[i + 2], hi = b[i + 3];
                bool looksWord = (hi == 0x00) || (lo <= 0x20 && hi <= 0x04);
                return looksWord ? (i + 4) : (i + 3);
            }

            if (i + 2 < endAbs) return (i + 3);
            return iAbs;
        }

        
        private static void SelectPreviewWindow(byte[] b, int start, int len, out int winStartAbs, out int winLen)
        {
            winStartAbs = start; winLen = 0;
            int relAnchor, anchorLen;
            FindLongestAsciiRun(b, start, len, out relAnchor, out anchorLen);

            int absStart;
            if (relAnchor >= 0)
            {
                absStart = start + relAnchor;
            }
            else
            {
                
                int end = Math.Min(b.Length, start + len);
                absStart = end; 
                for (int i = start; i < end; i++)
                {
                    if (IsPrintable(b[i])) { absStart = i; break; }
                }
                if (absStart >= end)
                {
                    
                    winStartAbs = start; winLen = 0;
                    return;
                }
            }

            int endAbsBound = Math.Min(b.Length, start + len);
            int i2 = absStart;

            
            while (i2 < endAbsBound)
            {
                if (IsPrintable(b[i2]))
                {
                    i2++;
                    continue;
                }

                if (b[i2] == 0x90)
                {
                    
                    if (i2 + 1 < endAbsBound && b[i2 + 1] == 0x00)
                        break;

                    int next = ConsumeWildcardBlock(b, start, len, i2);
                    if (next == i2) break; 
                    i2 = next;
                    continue;
                }

                
                break;
            }

            winStartAbs = absStart;
            winLen = Math.Max(0, i2 - absStart);
        }

        
        private static void LogPayloadPreviews(byte[] buf, int start, int len, string pad)
        {
            if (len <= 0)
            {
                Logger.Info(pad + "(no data)");
                return;
            }

            
            Logger.Verbose(string.Format(pad + "HexFull({0}): {1}", len, Hex(buf, start, len)));
            if (len > 64)
            {
                string head = Hex(buf, start, 64);
                Logger.Info(string.Format(pad + "HexHead(64): {0} ...", head));
            }

            
            int winStartAbs, winLen;
            SelectPreviewWindow(buf, start, len, out winStartAbs, out winLen);

            if (winLen <= 0)
            {
                Logger.Info(string.Format(pad + "Ascii/Wild @+0x{0:X}: ", start));
                Logger.Info(string.Format(pad + "ASCII-only @+0x{0:X}: \"\"", start));
                Logger.Info(string.Format(pad + "YARA-hex  @+0x{0:X}: ", start));
                return;
            }

            var slice = new byte[winLen];
            Buffer.BlockCopy(buf, winStartAbs, slice, 0, winLen);

            
            int consumed;
            bool hadTerm;
            var tokens = WildcardPattern.Tokenize(slice, 0, slice.Length, out consumed, out hadTerm);

            string human = WildcardPattern.RenderHuman(tokens, Encoding.ASCII) ?? string.Empty;
            string yara = WildcardPattern.RenderYaraHex(tokens) ?? string.Empty;

            
            var sbAscii = new StringBuilder();
            for (int k = 0; k < tokens.Count; k++)
            {
                var lit = tokens[k] as TokLiteral;
                if (lit != null && lit.Bytes != null && lit.Bytes.Length > 0)
                    sbAscii.Append(Encoding.ASCII.GetString(lit.Bytes));
            }

            Logger.Info(string.Format(pad + "Ascii/Wild @+0x{0:X}: {1}", winStartAbs, human));
            Logger.Info(string.Format(pad + "ASCII-only @+0x{0:X}: \"{1}\"", winStartAbs, sbAscii.ToString()));
            Logger.Info(string.Format(pad + "YARA-hex  @+0x{0:X}: {1}", winStartAbs, yara));
        }

        private static string DescribeLogic(byte type)
        {
            switch (type)
            {
                case 0x01: return "AND";
                case 0x02: return "LEAF";
                case 0x03: return "OR";
                case 0x04: return "NOT";
                case 0x0E: return "MATCH";
                case 0x1D: return "LOGIC_NODE";
                default: return string.Format("TYPE_0x{0:X2}", type);
            }
        }
    }
}