// DefenderRuleParser
// Author: Andrea Cristaldi 2025 - https://github.com/andreacristaldi/DefenderRuleParser
// This project is licensed under the Apache 2.0 License.
/*
 * Summary: Registry key/value extractor from signature blobs.
 * Origin: dump-driven; UTF-16LE registry paths and value names parsed heuristically.
 * Role: Outputs normalized HK* paths for triage and export.
 */
using DefenderRuleParser2;
using DefenderRuleParser2.Models;
using DefenderRuleParser2.Parsers.Wildcards; 
using System;
using System.Collections.Generic;
using System.IO;
using System.Text;

namespace DefenderRuleParser2.Parsers
{
    
   
    public class RegKeyParser : ISignatureParser
    {
        public void Parse(BinaryReader reader, int size, uint threatId)
        {
            long offset = reader.BaseStream.Position;

            try
            {
                byte[] buffer = reader.ReadBytes(size);
                if (buffer == null || buffer.Length < 4)
                {
                    Logger.Warn(string.Format("[REGKEY] Short payload ({0}) @0x{1:X}", (buffer == null ? 0 : buffer.Length), offset));
                    return;
                }

                using (var ms = new MemoryStream(buffer))
                using (var br = new BinaryReader(ms))
                {
                    ushort flags = br.ReadUInt16();
                    ushort threshold = br.ReadUInt16();

                    
                    byte[] keyBytes = ReadRest(br);

                   
                    bool looksWild = LooksLikeWildcardPattern(keyBytes);

                    string humanOut;
                    if (looksWild)
                    {
                        
                        int consumed;
                        bool hadTerm;
                        var tokens = WildcardPattern.Tokenize(keyBytes, 0, keyBytes.Length, out consumed, out hadTerm);

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

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

                        Logger.Info(string.Format("[REGKEY] Threat ID: {0}, Flags=0x{1:X4}, Threshold={2} (wildcarded)", threatId, flags, threshold));
                        if (!string.IsNullOrEmpty(asciiOnly))
                            Logger.Info("  · ASCII-only : " + asciiOnly);
                        if (!string.IsNullOrEmpty(human))
                            Logger.Info("  · WILD/Human : " + human);
                        if (!string.IsNullOrEmpty(yara))
                            Logger.Info("  · WILD/YARA  : " + yara);

                        humanOut = !string.IsNullOrEmpty(human) ? human : asciiOnly;
                    }
                    else
                    {
                        
                        string regKey = DecodeBestEffort(keyBytes);
                        Logger.Info(string.Format("[REGKEY] Threat ID: {0}, Flags=0x{1:X4}, Threshold={2}", threatId, flags, threshold));
                        Logger.Info("  > Key:   " + regKey);
                        Logger.Info("  > Flags: " + DescribeFlags(flags));
                        humanOut = regKey;
                    }

                    if (!string.IsNullOrWhiteSpace(humanOut) &&
                        ThreatDatabase.TryGetThreat(threatId, out var threat))
                    {
                        threat.Signatures.Add(new SignatureEntry
                        {
                            Type = "SIGNATURE_TYPE_REGKEY",
                            Offset = offset,
                            Pattern = new List<string> { humanOut },
                            Parsed = true,
                            
                            ConditionType = (threshold > 1) ? "MIN_MATCHES" : "PRESENT",
                            ConditionValue = Math.Max(1, (int)threshold)
                        });
                    }
                }
            }
            catch (Exception ex)
            {
                Logger.Error(string.Format("[!] REGKEY Error parsing at 0x{0:X}: {1}", offset, ex.Message));
            }
            finally
            {
                reader.BaseStream.Seek(offset + size, SeekOrigin.Begin);
            }
        }


        private static byte[] ReadRest(BinaryReader br)
        {
            int left = (int)(br.BaseStream.Length - br.BaseStream.Position);
            if (left <= 0) return new byte[0];
            return br.ReadBytes(left);
        }

        
        private static bool LooksLikeWildcardPattern(byte[] data)
        {
            if (data == null || data.Length < 2) return false;
            for (int i = 0; i < data.Length - 1; i++)
            {
                if (data[i] != 0x90) continue;
                byte next = data[i + 1];
                if (next == 0x00 || next == 0x90 || next < 0x32) return true;
            }
            return false;
        }

      
        private static string DecodeBestEffort(byte[] keyBytes)
        {
            if (keyBytes == null || keyBytes.Length == 0) return string.Empty;

            
            int nullOdd = 0;
            for (int i = 1; i < keyBytes.Length; i += 2)
                if (keyBytes[i] == 0x00) nullOdd++;

            bool looksUtf16 = (keyBytes.Length >= 4) && (nullOdd >= keyBytes.Length / 4);

            string s = looksUtf16
                ? Encoding.Unicode.GetString(keyBytes)
                : Encoding.UTF8.GetString(keyBytes);

            
            s = s.TrimEnd('\0').Trim();
            return s;
        }

        private string DescribeFlags(ushort flags)
        {
            var meanings = new List<string>(4);
            if ((flags & 0x0001) != 0) meanings.Add("Case Sensitive");
            if ((flags & 0x0002) != 0) meanings.Add("Match");
            if ((flags & 0x0004) != 0) meanings.Add("Optional");
            if (meanings.Count == 0) return "None";
            return string.Join(", ", meanings.ToArray());
        }
    }
}