unreal-sdk/src/sdk/repr.rs

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::v2_types::{
    traits::{UObjectNonConst, UObjectTrait},
    EFunctionFlags, EPropertyFlags, UObject,
};

impl UObject {
    pub fn as_package_ref(&self) -> Option<PackageRef> {
        if self.is_package_object() {
            Some(PackageRef(*self.internal_index()))
        } else {
            None
        }
    }

    pub fn object_ref(&self) -> ObjectRef {
        ObjectRef {
            package: PackageRef(*self.package_object().internal_index()),
            object: *self.internal_index(),
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Serialize, Deserialize)]
pub struct PackageRef(u32);

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Serialize, Deserialize)]
pub struct ObjectRef {
    pub package: PackageRef,
    pub object: u32,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Sdk {
    packages: BTreeMap<PackageRef, ProcessedPackage>,
}

/// A package represents some group of packages that are related to another,
/// examples in ARK are `Engine`, `ShooterGame`, `SlateCore` or
/// `Buff_Companion_HLNA`.
#[derive(Debug)]
pub struct Package {
    /// the package object is the root object of the package.
    pub package_object: UObject,
    /// each package may contain other objects and classes which can be
    /// referenced by an `ObjectRef`, which is a unique identifier for each
    /// object.
    pub children: Vec<UObject>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ProcessedPackage {
    /// the package object ref.
    pub package_object: PackageRef,
    /// all types extracted from this package referenced by their `ObjectRef`.
    pub types: BTreeMap<ObjectRef, UnrealType>,
    /// All other packages that types in this package depend on directly.
    pub dependencies: Vec<PackageRef>,
}

#[derive(Debug, Serialize, Deserialize)]
pub enum UnrealType {
    Class(Class),
    Struct(Class),
    Actor(Class),
    Enum(Enum),
}

impl UnrealType {
    pub fn obj_ref(&self) -> ObjectRef {
        match self {
            UnrealType::Class(obj) => obj.obj_ref,
            UnrealType::Struct(obj) => obj.obj_ref,
            UnrealType::Actor(obj) => obj.obj_ref,
            UnrealType::Enum(obj) => obj.obj_ref,
        }
    }

    pub fn get_dependent_types(&self) -> Vec<ObjectRef> {
        match self {
            UnrealType::Class(obj) | UnrealType::Struct(obj) | UnrealType::Actor(obj) => {
                obj.get_dependent_types()
            }
            UnrealType::Enum(_) => vec![],
        }
    }
}

#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum StructKind {
    Object,
    Actor,
    Struct,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Enum {
    pub obj_ref: ObjectRef,
    pub name: String,
    pub full_name: String,
    pub values: BTreeMap<u32, String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Class {
    pub obj_ref: ObjectRef,
    pub kind: StructKind,
    pub size: u32,
    pub name: String,
    pub full_name: String,
    pub super_class: Option<ObjectRef>,
    pub properties_size: u32,
    pub min_alignment: u32,
    pub fields: Vec<ClassField>,
    pub methods: Vec<ClassMethod>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ClassField {
    pub offset: u32,
    pub size: u32,
    pub name: String,
    pub flags: EPropertyFlags,
    pub ty: Type,
}

impl ClassField {
    pub fn is_return_param(&self) -> bool {
        self.flags.contains(EPropertyFlags::ReturnParm)
    }

    pub fn is_param(&self) -> bool {
        self.flags.contains(EPropertyFlags::Parm)
    }

    pub fn is_const_param(&self) -> bool {
        self.flags.contains(EPropertyFlags::ConstParm)
    }

    pub fn is_out_param(&self) -> bool {
        !self.is_const_param() && self.flags.contains(EPropertyFlags::OutParm)
    }

    pub fn is_rep(&self) -> bool {
        !self.flags.contains(EPropertyFlags::RepSkip)
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ClassMethod {
    pub name: String,
    pub full_name: String,
    pub flags: EFunctionFlags,
    pub parameters: Vec<ClassField>,
}

impl ClassMethod {
    pub fn is_static(&self) -> bool {
        self.flags.contains(EFunctionFlags::Static)
    }
    pub fn is_native(&self) -> bool {
        self.flags.contains(EFunctionFlags::Native)
    }
    /// this function is replicated to the server, might be exploitable.
    pub fn is_net_server(&self) -> bool {
        self.flags.contains(EFunctionFlags::NetServer)
    }
    pub fn out_params(&self) -> Vec<&ClassField> {
        self.parameters
            .iter()
            .filter(|&param| param.is_out_param() || param.is_return_param())
            .collect::<Vec<_>>()
    }

    pub fn in_params(&self) -> Vec<&ClassField> {
        self.parameters
            .iter()
            .filter(|param| param.is_param() && !(param.is_out_param() || param.is_return_param()))
            .collect::<Vec<_>>()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum PrimitiveType {
    Bool {
        byte_mask: u8,
        field_mask: u8,
        byte_offset: u8,
        field_size: u8,
    },
    U8,
    U16,
    U32,
    U64,
    I8,
    I16,
    I32,
    I64,
    F32,
    F64,
    Custom(String),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Type {
    Ptr(Box<Type>),
    Ref(Box<Type>),
    WeakPtr(ObjectRef),
    SoftPtr(ObjectRef),
    LazyPtr(ObjectRef),
    AssetPtr(ObjectRef),
    Array(Box<Type>),
    Primitive(PrimitiveType),
    RawArray {
        ty: Box<Type>,
        len: u32,
    },
    Name,
    String,
    Text,
    Enum {
        underlying: Box<Type>,
        enum_type: ObjectRef,
    },
    Class(ObjectRef),
    Struct(ObjectRef),
}

impl Type {
    pub fn dependent_type(&self) -> Option<ObjectRef> {
        match self {
            Type::Ptr(ty) | Type::Ref(ty) | Type::Array(ty) | Type::RawArray { ty, .. } => {
                ty.dependent_type()
            }
            Type::WeakPtr(ty) | Type::SoftPtr(ty) | Type::LazyPtr(ty) | Type::AssetPtr(ty) => {
                Some(*ty)
            }
            Type::Enum { underlying, .. } => underlying.dependent_type(),
            Type::Class(ty) | Type::Struct(ty) => Some(*ty),
            _ => None,
        }
    }
}