obfuz/Editor/Rename/SymbolRename.cs

using dnlib.DotNet;
using Obfuz.Rename;
using System;
using System.Collections;
using System.Collections.Generic;
using System.Data.OleDb;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using UnityEditor.SceneManagement;
using UnityEngine;
using UnityEngine.Assertions;

namespace Obfuz
{


    public class SymbolRename
    {
        private readonly AssemblyCache _assemblyCache;
        private readonly List<ObfuzAssemblyInfo> _obfuzAssemblies;
        private readonly HashSet<ModuleDef> _obfuscatedModules = new HashSet<ModuleDef>();
        private readonly IRenamePolicy _renamePolicy;
        private readonly INameMaker _nameMaker;
        private readonly Dictionary<ModuleDef, List<CustomAttributeInfo>> _customAttributeArgumentsWithTypeByMods = new Dictionary<ModuleDef, List<CustomAttributeInfo>>();
        private readonly RenameRecordMap _renameRecordMap;
        private readonly VirtualMethodGroupCalculator _virtualMethodGroupCalculator = new VirtualMethodGroupCalculator();

        class CustomAttributeInfo
        {
            public CustomAttributeCollection customAttributes;
            public int index;
            public List<CAArgument> arguments;
            public List<CANamedArgument> namedArguments;
        }
        public SymbolRename(ObfuscatorContext ctx)
        {
            _assemblyCache = ctx.assemblyCache;
            _obfuzAssemblies = ctx.assemblies;
            _renamePolicy = ctx.renamePolicy;
            _nameMaker = ctx.nameMaker;

            foreach (var mod in ctx.assemblies)
            {
                _obfuscatedModules.Add(mod.module);
            }
            BuildCustomAttributeArguments();

            _renameRecordMap = new RenameRecordMap(Path.Combine(ctx.outputDir, "mapping.xml"));
        }

        private void CollectCArgumentWithTypeOf(IHasCustomAttribute meta, List<CustomAttributeInfo> customAttributes)
        {
            int index = 0;
            foreach (CustomAttribute ca in meta.CustomAttributes)
            {
                List<CAArgument> arguments = null;
                if (ca.ConstructorArguments.Any(a => MetaUtil.MayRenameCustomDataType(a.Type.ElementType)))
                {
                    arguments = ca.ConstructorArguments.ToList();
                }
                List<CANamedArgument> namedArguments = null;
                if (ca.NamedArguments.Any(a => MetaUtil.MayRenameCustomDataType(a.Type.ElementType)))
                {
                    namedArguments = ca.NamedArguments.ToList();
                }
                if (arguments != null | namedArguments != null)
                {
                    customAttributes.Add(new CustomAttributeInfo
                    {
                        customAttributes = meta.CustomAttributes,
                        index = index,
                        arguments = arguments,
                        namedArguments = namedArguments
                    });
                }
                ++index;
            }
        }

        private void BuildCustomAttributeArguments()
        {
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                var customAttributes = new List<CustomAttributeInfo>();
                CollectCArgumentWithTypeOf(ass.module, customAttributes);
                foreach (TypeDef type in ass.module.GetTypes())
                {
                    CollectCArgumentWithTypeOf(type, customAttributes);
                    foreach (FieldDef field in type.Fields)
                    {
                        CollectCArgumentWithTypeOf(field, customAttributes);
                    }
                    foreach (MethodDef method in type.Methods)
                    {
                        CollectCArgumentWithTypeOf(method, customAttributes);
                        foreach (Parameter param in method.Parameters)
                        {
                            if (param.ParamDef != null)
                            {
                                CollectCArgumentWithTypeOf(param.ParamDef, customAttributes);
                            }
                        }
                    }
                    foreach (PropertyDef property in type.Properties)
                    {
                        CollectCArgumentWithTypeOf(property, customAttributes);
                    }
                    foreach (EventDef eventDef in type.Events)
                    {
                        CollectCArgumentWithTypeOf (eventDef, customAttributes);
                    }
                }

                _customAttributeArgumentsWithTypeByMods.Add(ass.module, customAttributes);
            }
        }

        public void Process()
        {
            _renameRecordMap.Init(_obfuzAssemblies, _nameMaker);
            RenameModules();
            RenameTypes();
            RenameFields();
            RenameMethods();
            RenameProperties();
            RenameEvents();
        }

        private List<ObfuzAssemblyInfo> GetReferenceMeAssemblies(ModuleDef mod)
        {
            return _obfuzAssemblies.Find(ass => ass.module == mod).referenceMeAssemblies;
        }

        private void RenameModules()
        {
            Debug.Log("Rename Modules begin");
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                if (_renamePolicy.NeedRename(ass.module))
                {
                    Rename(ass.module);
                }
            }
            Debug.Log("Rename Modules end");
        }


        class RefTypeDefMetas
        {
            public readonly List<TypeRef> typeRefs = new List<TypeRef>();

            public readonly List<CustomAttribute> customAttributes = new List<CustomAttribute>();
        }

        private void BuildRefTypeDefMetasMap(Dictionary<TypeDef, RefTypeDefMetas> refTypeDefMetasMap)
        {
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (TypeRef typeRef in ass.module.GetTypeRefs())
                {
                    TypeDef typeDef = typeRef.ResolveThrow();
                    if (!refTypeDefMetasMap.TryGetValue(typeDef, out var typeDefMetas))
                    {
                        typeDefMetas = new RefTypeDefMetas();
                        refTypeDefMetasMap.Add(typeDef, typeDefMetas);
                    }
                    typeDefMetas.typeRefs.Add(typeRef);
                }
            }

            foreach (CustomAttributeInfo cai in _customAttributeArgumentsWithTypeByMods.Values.SelectMany(cas => cas))
            {
                CustomAttribute ca = cai.customAttributes[cai.index];
                TypeDef typeDef = MetaUtil.GetTypeDefOrGenericTypeBaseThrowException(ca.Constructor.DeclaringType);
                if (!refTypeDefMetasMap.TryGetValue(typeDef, out var typeDefMetas))
                {
                    typeDefMetas = new RefTypeDefMetas();
                    refTypeDefMetasMap.Add(typeDef, typeDefMetas);
                }
                typeDefMetas.customAttributes.Add(ca);
            }
        }

        private void RetargetTypeRefInCustomAttributes()
        {
            foreach (CustomAttributeInfo cai in _customAttributeArgumentsWithTypeByMods.Values.SelectMany(cas => cas))
            {
                CustomAttribute ca = cai.customAttributes[cai.index];
                bool anyChange = false;
                if (cai.arguments != null)
                {
                    for (int i = 0; i < cai.arguments.Count; i++)
                    {
                        CAArgument oldArg = cai.arguments[i];
                        if (MetaUtil.TryRetargetTypeRefInArgument(oldArg, out CAArgument newArg))
                        {
                            anyChange = true;
                            cai.arguments[i] = newArg;
                        }
                    }
                }
                if (cai.namedArguments != null)
                {
                    for (int i = 0; i < cai.namedArguments.Count; i++)
                    {
                        if (MetaUtil.TryRetargetTypeRefInNamedArgument(cai.namedArguments[i]))
                        {
                            anyChange = true;
                        }
                    }
                }
                if (anyChange)
                {
                    cai.customAttributes[cai.index] = new CustomAttribute(ca.Constructor,
                        cai.arguments != null ? cai.arguments : ca.ConstructorArguments,
                        cai.namedArguments != null ? cai.namedArguments : ca.NamedArguments);
                }
            }
        }

        private readonly Dictionary<TypeDef, RefTypeDefMetas> _refTypeRefMetasMap = new Dictionary<TypeDef, RefTypeDefMetas>();

        private void RenameTypes()
        {
            Debug.Log("RenameTypes begin");

            RetargetTypeRefInCustomAttributes();

            BuildRefTypeDefMetasMap(_refTypeRefMetasMap);
            _assemblyCache.EnableTypeDefCache = false;

            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (TypeDef type in ass.module.GetTypes())
                {
                    if (_renamePolicy.NeedRename(type))
                    {
                        Rename(type, _refTypeRefMetasMap.GetValueOrDefault(type));
                    }
                }
            }

            // clean cache
            _assemblyCache.EnableTypeDefCache = true;
            Debug.Log("Rename Types end");
        }


        class RefFieldMetas
        {
            public readonly List<MemberRef> fieldRefs = new List<MemberRef>();
            public readonly List<CustomAttribute> customAttributes = new List<CustomAttribute>();
        }


        private void BuildHierarchyFields(TypeDef type, List<FieldDef> fields)
        {
            while (type != null)
            {
                fields.AddRange(type.Fields);
                type = MetaUtil.GetBaseTypeDef(type);
            }
        }

        private void BuildRefFieldMetasMap(Dictionary<FieldDef, RefFieldMetas> refFieldMetasMap)
        {
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (MemberRef memberRef in ass.module.GetMemberRefs())
                {
                    if (!memberRef.IsFieldRef)
                    {
                        continue;
                    }
                     
                    IMemberRefParent parent = memberRef.Class;
                    TypeDef parentTypeDef = MetaUtil.GetMemberRefTypeDefParentOrNull(parent);
                    if (parentTypeDef == null)
                    {
                        continue;
                    }
                    foreach (FieldDef field in parentTypeDef.Fields)
                    {
                        if (field.Name == memberRef.Name && TypeEqualityComparer.Instance.Equals(field.FieldSig.Type, memberRef.FieldSig.Type))
                        {
                            if (!refFieldMetasMap.TryGetValue(field, out var fieldMetas))
                            {
                                fieldMetas = new RefFieldMetas();
                                refFieldMetasMap.Add(field, fieldMetas);
                            }
                            fieldMetas.fieldRefs.Add(memberRef);
                            break;
                        }
                    }
                }
            }
            foreach (var e in _refTypeRefMetasMap)
            {
                TypeDef typeDef = e.Key;
                var hierarchyFields = new List<FieldDef>();
                BuildHierarchyFields(typeDef, hierarchyFields);
                RefTypeDefMetas typeDefMetas = e.Value;
                foreach (CustomAttribute ca in typeDefMetas.customAttributes)
                {
                    foreach (var arg in ca.NamedArguments)
                    {
                        if (arg.IsProperty)
                        {
                            continue;
                        }
                        foreach (FieldDef field in hierarchyFields)
                        {
                            // FIXME. field of Generic Base Type may not be same
                            if (field.Name == arg.Name && TypeEqualityComparer.Instance.Equals(field.FieldType, arg.Type))
                            {
                                if (!refFieldMetasMap.TryGetValue(field, out var fieldMetas))
                                {
                                    fieldMetas = new RefFieldMetas();
                                    refFieldMetasMap.Add(field, fieldMetas);
                                }
                                fieldMetas.customAttributes.Add(ca);
                                break;
                            }
                        }
                    }
                }
            }
        }

        private void RenameFields()
        {
            Debug.Log("Rename fields begin");
            var refFieldMetasMap = new Dictionary<FieldDef, RefFieldMetas>();
            BuildRefFieldMetasMap(refFieldMetasMap);

            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (TypeDef type in ass.module.GetTypes())
                {
                    foreach (FieldDef field in type.Fields)
                    {
                        if (_renamePolicy.NeedRename(field))
                        {
                            Rename(field, refFieldMetasMap.GetValueOrDefault(field));
                        }
                    }
                }
            }
            Debug.Log("Rename fields end");
        }

        class RefMethodMetas
        {
            public readonly List<MemberRef> memberRefs = new List<MemberRef>();
        }

        private void BuildRefMethodMetasMap(Dictionary<MethodDef, RefMethodMetas> refMethodMetasMap)
        {
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (MemberRef memberRef in ass.module.GetMemberRefs())
                {
                    if (!memberRef.IsMethodRef)
                    {
                        continue;
                    }

                    IMemberRefParent parent = memberRef.Class;
                    TypeDef parentTypeDef = MetaUtil.GetMemberRefTypeDefParentOrNull(parent);
                    if (parentTypeDef == null)
                    {
                        continue;
                    }
                    foreach (MethodDef method in parentTypeDef.Methods)
                    {
                        if (method.Name == memberRef.Name && new SigComparer(default).Equals(method.MethodSig, memberRef.MethodSig))
                        {
                            if (!refMethodMetasMap.TryGetValue(method, out var refMethodMetas))
                            {
                                refMethodMetas = new RefMethodMetas();
                                refMethodMetasMap.Add(method, refMethodMetas);
                            }
                            refMethodMetas.memberRefs.Add(memberRef);
                            break;
                        }
                    }
                }
            }
        }

        private void RenameMethods()
        {
            Debug.Log("Rename methods begin");
            Debug.Log("Rename not virtual methods begin");
            var virtualMethods = new List<MethodDef>();
            var refMethodMetasMap = new Dictionary<MethodDef, RefMethodMetas>();
            BuildRefMethodMetasMap(refMethodMetasMap);
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (TypeDef type in ass.module.GetTypes())
                {
                    _virtualMethodGroupCalculator.CalculateType(type);
                    foreach (MethodDef method in type.Methods)
                    {
                        if (method.IsVirtual)
                        {
                            virtualMethods.Add(method);
                            continue;
                        }
                        if (_renamePolicy.NeedRename(method))
                        {
                            Rename(method, refMethodMetasMap.GetValueOrDefault(method));
                        }
                    }
                }
            }
            Debug.Log("Rename not virtual methods end");


            Debug.Log("Rename virtual methods begin");
            var visitedVirtualMethods = new HashSet<MethodDef>();
            var groupNeedRenames = new Dictionary<VirtualMethodGroup, bool>();
            foreach (var method in virtualMethods)
            {
                if (!visitedVirtualMethods.Add(method))
                {
                    continue;
                }
                VirtualMethodGroup group = _virtualMethodGroupCalculator.GetMethodGroup(method);
                if (!groupNeedRenames.TryGetValue(group, out var needRename))
                {
                    needRename = group.methods.All(m => _obfuscatedModules.Contains(m.DeclaringType.Module) && _renamePolicy.NeedRename(m));
                    groupNeedRenames.Add(group, needRename);
                    if (needRename)
                    {
                        _renameRecordMap.InitAndAddRename(group, _nameMaker.GetNewName(method, method.Name));
                    }
                }
                if (!needRename)
                {
                    continue;
                }
                if (_renameRecordMap.TryGetRename(group, out var newName))
                {
                    Rename(method, refMethodMetasMap.GetValueOrDefault(method), newName);
                }
                else
                {
                    throw new Exception($"group:{group} method:{method} not found in rename record map");
                }
            }
            Debug.Log("Rename virtual methods end");
            Debug.Log("Rename methods end");
        }

        class RefPropertyMetas
        {
            public List<CustomAttribute> customAttributes = new List<CustomAttribute>();
        }

        private void BuildHierarchyProperties(TypeDef type, List<PropertyDef> properties)
        {
            while (type != null)
            {
                properties.AddRange(type.Properties);
                type = MetaUtil.GetBaseTypeDef(type);
            }
        }

        private void BuildRefPropertyMetasMap(Dictionary<PropertyDef, RefPropertyMetas> refPropertyMetasMap)
        {
            foreach (var e in _refTypeRefMetasMap)
            {
                TypeDef typeDef = e.Key;
                var hierarchyProperties = new List<PropertyDef>();
                BuildHierarchyProperties(typeDef, hierarchyProperties);
                RefTypeDefMetas typeDefMetas = e.Value;
                foreach (CustomAttribute ca in typeDefMetas.customAttributes)
                {
                    foreach (var arg in ca.NamedArguments)
                    {
                        if (arg.IsField)
                        {
                            continue;
                        }
                        foreach (PropertyDef field in hierarchyProperties)
                        {
                            // FIXME. field of Generic Base Type may not be same
                            if (field.Name == arg.Name && TypeEqualityComparer.Instance.Equals(arg.Type, field.PropertySig.RetType))
                            {
                                if (!refPropertyMetasMap.TryGetValue(field, out var fieldMetas))
                                {
                                    fieldMetas = new RefPropertyMetas();
                                    refPropertyMetasMap.Add(field, fieldMetas);
                                }
                                fieldMetas.customAttributes.Add(ca);
                                break;
                            }
                        }
                    }
                }
            }
        }

        private void RenameProperties()
        {
            Debug.Log("Rename properties begin");
            var refPropertyMetasMap = new Dictionary<PropertyDef, RefPropertyMetas>();
            BuildRefPropertyMetasMap(refPropertyMetasMap);
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (TypeDef type in ass.module.GetTypes())
                {
                    foreach (PropertyDef property in type.Properties)
                    {
                        if (_renamePolicy.NeedRename(property))
                        {
                            Rename(property, refPropertyMetasMap.GetValueOrDefault(property));
                        }
                    }
                }
            }
            Debug.Log("Rename properties end");
        }

        private void RenameEvents()
        {
            Debug.Log("Rename events begin");
            foreach (ObfuzAssemblyInfo ass in _obfuzAssemblies)
            {
                foreach (TypeDef type in ass.module.GetTypes())
                {
                    foreach (EventDef eventDef in type.Events)
                    {
                        if (_renamePolicy.NeedRename(eventDef))
                        {
                            Rename(eventDef);
                        }
                    }
                }
            }
            Debug.Log("Rename events begin");
        }

        private void Rename(ModuleDefMD mod)
        {
            string oldName = mod.Assembly.Name;
            string newName = _nameMaker.GetNewName(mod, oldName);
            _renameRecordMap.AddRename(mod, newName);
            mod.Assembly.Name = newName;
            mod.Name = $"{newName}.dll";
            //Debug.Log($"rename module. oldName:{oldName} newName:{newName}");
            foreach (ObfuzAssemblyInfo ass in GetReferenceMeAssemblies(mod))
            {
                foreach (AssemblyRef assRef in ass.module.GetAssemblyRefs())
                {
                    if (assRef.Name == oldName)
                    {
                        _renameRecordMap.AddRename(mod, newName);
                        assRef.Name = newName;
                       // Debug.Log($"rename assembly:{ass.name}  ref oldName:{oldName} newName:{newName}");
                    }
                }
            }
        }

        private void Rename(TypeDef type, RefTypeDefMetas refTypeDefMeta)
        {
            string moduleName = MetaUtil.GetModuleNameWithoutExt(type.Module.Name);
            string oldFullName = type.FullName;
            string oldNamespace = type.Namespace;
            string newNamespace;
            if (string.IsNullOrEmpty(oldNamespace))
            {
                newNamespace = oldNamespace;
            }
            else
            {
                newNamespace = _nameMaker.GetNewNamespace(type, oldNamespace, true);
                type.Namespace = newNamespace;
            }

            string oldName = type.Name;
            string newName = _nameMaker.GetNewName(type, oldName);

            if (refTypeDefMeta != null)
            {
                foreach (TypeRef typeRef in refTypeDefMeta.typeRefs)
                {
                    Assert.AreEqual(typeRef.FullName, oldFullName);
                    Assert.IsTrue(typeRef.DefinitionAssembly.Name == moduleName);
                    if (!string.IsNullOrEmpty(oldNamespace))
                    {
                        typeRef.Namespace = newNamespace;
                    }
                    typeRef.Name = newName;
                    //Debug.Log($"rename assembly:{typeRef.Module.Name} reference {oldFullName} => {typeRef.FullName}");
                }
            }
            type.Name = newName;
            string newFullName = type.FullName;
            _renameRecordMap.AddRename(type, newFullName);
            //Debug.Log($"rename typedef. assembly:{type.Module.Name} oldName:{oldFullName} => newName:{newFullName}");
        }

        private void Rename(FieldDef field, RefFieldMetas fieldMetas)
        {
            string oldName = field.Name;
            string newName = _nameMaker.GetNewName(field, oldName);
            if (fieldMetas != null)
            {
                foreach (var memberRef in fieldMetas.fieldRefs)
                {
                    memberRef.Name = newName;
                    //Debug.Log($"rename assembly:{memberRef.Module.Name} reference {field.FullName} => {memberRef.FullName}");
                }
                foreach (var ca in fieldMetas.customAttributes)
                {
                    foreach (var arg in ca.NamedArguments)
                    {
                        if (arg.Name == oldName)
                        {
                            arg.Name = newName;
                        }
                    }
                }
            }
            //Debug.Log($"rename field. {field} => {newName}");
            _renameRecordMap.AddRename(field, newName);
            field.Name = newName;

        }

        private void Rename(MethodDef method, RefMethodMetas refMethodMetas)
        {
            string oldName = method.Name;
            string newName = _nameMaker.GetNewName(method, oldName);
            Rename(method, refMethodMetas, newName);
        }

        private void Rename(MethodDef method, RefMethodMetas refMethodMetas, string newName)
        {

            ModuleDefMD mod = (ModuleDefMD)method.DeclaringType.Module;
            RenameMethodParams(method);
            RenameMethodBody(method);
            if (refMethodMetas != null)
            {
                foreach (MemberRef memberRef in refMethodMetas.memberRefs)
                {
                    string oldMethodFullName = memberRef.ToString();
                    memberRef.Name = newName;
                    //Debug.Log($"rename assembly:{memberRef.Module.Name} method:{oldMethodFullName} => {memberRef}");
                }
            }
            _renameRecordMap.AddRename(method, newName);
            method.Name = newName;
        }

        private void RenameMethodBody(MethodDef method)
        {
            if (method.Body == null)
            {
                return;
            }
        }

        private void RenameMethodParams(MethodDef method)
        {
            foreach (Parameter param in method.Parameters)
            {
                if (param.ParamDef != null)
                {
                    Rename(param.ParamDef);
                }
            }
        }

        private void Rename(ParamDef param)
        {
            if (_renamePolicy.NeedRename(param))
            {
                string newName = _nameMaker.GetNewName(param, param.Name);
                _renameRecordMap.AddRename(param, newName);
                param.Name = newName;
            }
        }

        private void Rename(EventDef eventDef)
        {
            string oldName = eventDef.Name;
            string newName = _nameMaker.GetNewName(eventDef, eventDef.Name);
            _renameRecordMap.AddRename(eventDef, newName);
            eventDef.Name = newName;
        }

        private void Rename(PropertyDef property, RefPropertyMetas refPropertyMetas)
        {
            string oldName = property.Name;
            string newName = _nameMaker.GetNewName(property, oldName);

            if (refPropertyMetas != null)
            {
                foreach (var ca in refPropertyMetas.customAttributes)
                {
                    foreach (var arg in ca.NamedArguments)
                    {
                        if (arg.Name == oldName)
                        {
                            arg.Name = newName;
                        }
                    }
                }
            }
            _renameRecordMap.AddRename(property, newName);
            property.Name = newName;
        }

        public void Save()
        {
            _renameRecordMap.WriteXmlMappingFile();
        }
    }
}