// DefenderRuleParser
// Author: Andrea Cristaldi 2025 - https://github.com/andreacristaldi/DefenderRuleParser
// This project is licensed under the Apache 2.0 License.
/*
 * Summary: Generic weighted-pattern parser (HSTR-like families, ARHSTR/PEHSTR/CMDHSTR/MACROHSTR).
 * Origin: dump-driven; header (threshold, count) and subrule layout derived from hex.
 * Role: Tokenizes wildcard-encoded bytes (0x90 scheme), falls back to string decoding, and preserves weights/thresholds.
 */
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 GenericWeightedPatternParser : ISignatureParser
    {
        private readonly string _signatureType;

        public GenericWeightedPatternParser(string signatureType)
        {
            _signatureType = signatureType;
        }

        public void Parse(BinaryReader reader, int size, uint threatId)
        {
            long offset = reader.BaseStream.Position;
            var patterns = new List<string>();
            var logic = new List<SubRuleLogic>();

            try
            {
                byte[] buffer = reader.ReadBytes(size);
                if (buffer == null || buffer.Length != size)
                {
                    Logger.Warn($"[{_signatureType}] Short read ({(buffer == null ? 0 : buffer.Length)}/{size}) @0x{offset:X}");
                    reader.BaseStream.Seek(offset + size, SeekOrigin.Begin);
                    return;
                }

                using (var ms = new MemoryStream(buffer))
                using (var br = new BinaryReader(ms))
                {
                    
                    if (br.BaseStream.Length < 7)
                    {
                        Logger.Warn($"[{_signatureType}] Size too small ({size}) @0x{offset:X}");
                        reader.BaseStream.Seek(offset + size, SeekOrigin.Begin);
                        return;
                    }

                    ushort unknown = br.ReadUInt16();
                    ushort threshold = br.ReadUInt16();
                    ushort subRuleCount = br.ReadUInt16();
                    byte ctrl = br.ReadByte(); 

                    Logger.Info($"[{_signatureType}] Threat ID: {threatId}, Threshold: {threshold}, SubRules: {subRuleCount}");

                    for (int i = 0; i < subRuleCount && br.BaseStream.Position < br.BaseStream.Length; i++)
                    {
                        long subStart = br.BaseStream.Position;

                        
                        if (subStart + 3 > br.BaseStream.Length)
                        {
                            Logger.Warn($"  ! SubRule header #{i + 1} truncated.");
                            break;
                        }

                        ushort weight = br.ReadUInt16();

                        if (br.BaseStream.Position >= br.BaseStream.Length)
                        {
                            Logger.Warn($"  ! SubRule header #{i + 1} missing type/len.");
                            break;
                        }

                        byte b1 = br.ReadByte();
                        byte typeByte = 0xFF;
                        byte ruleLen;

                        
                        if (b1 <= 0x02)
                        {
                            typeByte = b1;
                            if (br.BaseStream.Position >= br.BaseStream.Length)
                            {
                                Logger.Warn($"  ! SubRule header #{i + 1} missing length after type.");
                                break;
                            }
                            ruleLen = br.ReadByte();
                        }
                        else
                        {
                            ruleLen = b1;
                        }

                        
                        bool hadAsciiTag = false;
                        if (br.BaseStream.Position + 2 <= br.BaseStream.Length)
                        {
                            byte t0 = br.ReadByte();
                            byte t1 = br.ReadByte();
                            if (t0 == 0x80 && t1 == 0x01)
                            {
                                hadAsciiTag = true;
                            }
                            else
                            {
                                
                                br.BaseStream.Position -= 2;
                            }
                        }

                        if (br.BaseStream.Position + ruleLen > br.BaseStream.Length)
                        {
                            Logger.Warn($"  ! SubRule #{i + 1} truncated: size={ruleLen} beyond buffer.");
                            break;
                        }

                        byte[] patternBytes = br.ReadBytes(ruleLen);

                        
                        string renderedHuman = null;
                        string renderedYara = null;
                        bool usedTokenizer = false;

                        
                        if (!hadAsciiTag && LooksLikeWildcardPattern(patternBytes))
                        {
                            int consumed; bool hadTerm;
                            var tokens = WildcardPattern.Tokenize(patternBytes, 0, patternBytes.Length, out consumed, out hadTerm);
                            if (tokens != null && tokens.Count > 0)
                            {
                                renderedHuman = WildcardPattern.RenderHuman(tokens);
                                renderedYara = WildcardPattern.RenderYaraHex(tokens);
                                usedTokenizer = true;

                                Logger.Info($"  > SubRule #{i + 1}: Weight={weight}, WildTokenized=true, Len={ruleLen}, Term={(hadTerm ? "yes" : "no")}");
                                Logger.Info($"    · Human   : {renderedHuman}");
                                Logger.Info($"    · YARA hex: {renderedYara}");
                            }
                        }

                        
                        if (!usedTokenizer)
                        {
                            string decoded = IsLikelyUtf16(patternBytes)
                                ? Encoding.Unicode.GetString(patternBytes).TrimEnd('\0')
                                : Encoding.UTF8.GetString(patternBytes).TrimEnd('\0');

                            renderedHuman = decoded;

                            string suffix =
                                hadAsciiTag ? " (ascii-tag)" :
                                (typeByte <= 0x02 ? $" (type=0x{typeByte:X2})" : "");

                            Logger.Info($"  > SubRule #{i + 1}: Weight={weight}, Pattern=\"{decoded}\"{suffix}");
                        }

                        
                        patterns.Add(renderedHuman ?? string.Empty);

                        logic.Add(new SubRuleLogic
                        {
                            Pattern = renderedHuman ?? string.Empty,
                            Weight = weight,
                            Control = (typeByte <= 0x02) ? typeByte : (byte)0
                        });
                    }

                    if (ThreatDatabase.TryGetThreat(threatId, out var threat))
                    {
                        threat.Signatures.Add(new SignatureEntry
                        {
                            Type = _signatureType,
                            Offset = offset,
                            Pattern = patterns,
                            Parsed = true,
                            Logic = new SignatureLogic
                            {
                                Threshold = threshold,
                                SubRules = logic
                            }
                        });
                    }
                }
            }
            catch (Exception ex)
            {
                Logger.Error($"[!] {_signatureType} Error parsing at offset 0x{offset:X}: {ex.Message}");
            }
            finally
            {
                reader.BaseStream.Seek(offset + size, SeekOrigin.Begin);
            }
             

        }

        
        private bool LooksLikeWildcardPattern(byte[] data)
        {
            if (data == null || data.Length == 0) return false;

            for (int i = 0; i < data.Length; i++)
            {
                if (data[i] != 0x90) continue;

                if (i + 1 < data.Length)
                {
                    byte next = data[i + 1];
                    if (next == 0x00 || next == 0x90) return true; 
                    if (next < 0x32) return true;                   
                }
            }
            return false;
        }

       
        private bool IsLikelyUtf16(byte[] data)
        {
            if (data == null || data.Length < 2) return false;

            int nulls = 0;
            for (int i = 1; i < data.Length; i += 2)
            {
                if (data[i] == 0x00) nulls++;
            }
            return (nulls >= data.Length / 4);
        }
    }
}