analysis.hpp

//--------------------------------------------------------------------*- C++ -*-
// Kunai-static-analyzer: library for doing analysis of dalvik files
// @author Farenain <kunai.static.analysis@gmail.com>
// @author Ernesto Java <javaernesto@gmail.com>
//
// @file analysis.hpp
// @brief Here we will have all the classes useful for analyst, for example
// accessing to cross-references, and easy accessing data from the DEX file.
 
#ifndef KUNAI_DEX_ANALYSIS_ANALYSIS_HPP
#define KUNAI_DEX_ANALYSIS_ANALYSIS_HPP
 
#include "Kunai/DEX/parser/parser.hpp"
#include "Kunai/DEX/DVM/disassembler.hpp"
#include "Kunai/DEX/analysis/external_class.hpp"
#include "Kunai/DEX/DVM/dalvik_instructions.hpp"
#include "Kunai/Exceptions/analysis_exception.hpp"
 
#include <set>
#include <variant>
 
namespace KUNAI
{
    namespace DEX
    {
        // forward declaration for the xrefs
        class ClassAnalysis;
        class MethodAnalysis;
        class FieldAnalysis;
        class BasicBlocks;
 
        class DVMBasicBlock
        {
            std::vector<Instruction *> instructions;
 
            bool try_block = false;
            bool catch_block = false;
            bool start_block = false;
            bool end_block = false;
            DVMType * handler_type;
 
            std::string name;
        public:
            DVMBasicBlock() = default;
 
            size_t get_nb_instructions() const
            {
                return instructions.size();
            }
 
            void add_instruction(Instruction *instr)
            {
                instructions.push_back(instr);
            }
 
            const std::vector<Instruction *> &get_instructions() const
            {
                return instructions;
            }
 
            std::vector<Instruction *> &get_instructions()
            {
                return instructions;
            }
 
            Instruction *get_terminator()
            {
                if (instructions.size() > 0 && instructions.back()->is_terminator())
                    return instructions.back();
                return nullptr;
            }
 
            std::uint64_t get_first_address()
            {
                if (instructions.size() > 0)
                    return instructions[0]->get_address();
                throw exceptions::AnalysisException("get_first_address(): basic block has no instructions");
            }
 
            std::uint64_t get_last_address()
            {
                if (instructions.size() > 0)
                    return instructions.back()->get_address();
                throw exceptions::AnalysisException("get_last_address(): basic block has no instructions");
            }
 
            std::string& get_name()
            {
                if (!name.empty())
                    return name;
 
                std::stringstream s;
 
                if (start_block)
                    name = "BB.(start block)";
                else if (end_block)
                    name = "BB.(end block)";
                else
                {
                    s << "BB." << get_first_address()
                        << "-" << get_last_address();
                    name = s.str();
                }
 
                return name;
            }
 
            bool is_try_block() const
            {
                return try_block;
            }
 
            void set_try_block(bool try_block)
            {
                this->try_block = try_block;
            }
 
            bool is_catch_block() const
            {
                return catch_block;
            }
 
            void set_catch_block(bool catch_block)
            {
                this->catch_block = catch_block;
            }
 
            DVMType * get_handler_type()
            {
                return handler_type;
            }
 
            void set_handler_type(DVMType * handler)
            {
                handler_type = handler;
            }
 
            bool is_start_block() const
            {
                return start_block;
            }
 
            void set_start_block(bool start_block)
            {
                this->start_block = start_block;
            }
 
            bool is_end_block() const
            {
                return end_block;
            }
 
            void set_end_block(bool end_block)
            {
                this->end_block = end_block;
            }
        };
 
        class BasicBlocks
        {
        public:
            using connected_blocks_t = std::unordered_map<DVMBasicBlock *,
                                                          std::set<DVMBasicBlock *>>;
            using edges_t = std::vector<std::pair<DVMBasicBlock *, DVMBasicBlock *>>;
 
            enum node_type_t
            {
                JOIN_NODE = 0, // len(predecessors) > 1
                BRANCH_NODE,   // len(sucessors) > 1
                REGULAR_NODE,  // other cases
            };
 
        private:
            std::vector<DVMBasicBlock *> nodes;
 
            connected_blocks_t predecessors;
 
            connected_blocks_t sucessors;
 
            edges_t edges;
 
        public:
            BasicBlocks() = default;
 
            ~BasicBlocks()
            {
                // free memory
                for (auto p : nodes)
                    delete p;
                // clear de vector of nodes
                nodes.clear();
            }
 
            size_t get_number_of_basic_blocks() const
            {
                return nodes.size();
            }
 
            const connected_blocks_t& get_predecessors() const
            {
                return predecessors;
            }
 
            connected_blocks_t& get_predecessors()
            {
                return predecessors;
            }
 
            void add_predecessor(DVMBasicBlock *node, DVMBasicBlock *pred)
            {
                predecessors[node].insert(pred);
            }
 
            const connected_blocks_t& get_sucessors() const
            {
                return sucessors;
            }
 
            connected_blocks_t& get_sucessors()
            {
                return sucessors;
            }
 
            void add_sucessor(DVMBasicBlock *node, DVMBasicBlock *suc)
            {
                sucessors[node].insert(suc);
            }
 
            void add_node(DVMBasicBlock *node)
            {
                if (std::find(nodes.begin(), nodes.end(), node) == nodes.end())
                    nodes.push_back(node);
            }
 
            void add_edge(DVMBasicBlock *src, DVMBasicBlock *dst)
            {
                add_node(src);
                add_node(dst);
                // now insert the edge
                auto edge_pair = std::make_pair(src, dst);
                if (std::find(edges.begin(), edges.end(), edge_pair) == edges.end())
                    edges.push_back(edge_pair);
                add_sucessor(src, dst);
                add_predecessor(dst, src);
            }
 
            const edges_t& get_edges() const
            {
                return edges;
            }
 
            edges_t& get_edges()
            {
                return edges;
            }
 
            node_type_t get_node_type(DVMBasicBlock *node)
            {
                if (predecessors[node].size() > 1)
                    return JOIN_NODE;
                else if (sucessors[node].size() > 1)
                    return BRANCH_NODE;
                else
                    return REGULAR_NODE;
            }
 
            void remove_node(DVMBasicBlock *node);
 
            DVMBasicBlock *get_basic_block_by_idx(std::uint64_t idx);
 
            const std::vector<DVMBasicBlock *>& get_nodes() const
            {
                return nodes;
            }
 
            std::vector<DVMBasicBlock *>& get_nodes()
            {
                return nodes;
            }
        };
 
        class FieldAnalysis
        {
            EncodedField *field;
            std::string &name;
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> xrefread;
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> xrefwrite;
 
        public:
            FieldAnalysis(EncodedField *field)
                : field(field), name(field->get_field()->get_name())
            {
            }
 
            EncodedField *get_field()
            {
                return field;
            }
 
            std::string &get_name()
            {
                return name;
            }
 
            const std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xrefread() const
            {
                return xrefread;
            }
 
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xrefread()
            {
                return xrefread;
            }
 
            const std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xrefwrite() const
            {
                return xrefwrite;
            }
 
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xrefwrite()
            {
                return xrefwrite;
            }
 
            void add_xrefread(ClassAnalysis *c, MethodAnalysis *m, std::uint64_t offset)
            {
                xrefread.push_back(std::make_tuple(c, m, offset));
            }
 
            void add_xrefwrite(ClassAnalysis *c, MethodAnalysis *m, std::uint64_t offset)
            {
                xrefwrite.push_back(std::make_tuple(c, m, offset));
            }
        };
 
        class StringAnalysis
        {
            std::string &value;
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> xreffrom;
 
        public:
            StringAnalysis(std::string &value) : value(value)
            {
            }
 
            const std::string &get_value() const
            {
                return value;
            }
 
            const std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xreffrom() const
            {
                return xreffrom;
            }
 
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xreffrom()
            {
                return xreffrom;
            }
 
            void add_xreffrom(ClassAnalysis *c, MethodAnalysis *m, std::uint64_t offset)
            {
                xreffrom.push_back(std::make_tuple(c, m, offset));
            }
        };
 
        class MethodAnalysis
        {
        public:
            const std::vector<std::string> known_apis{
                "Landroid/", "Lcom/android/internal/util", "Ldalvik/", "Ljava/", "Ljavax/", "Lorg/apache/",
                "Lorg/json/", "Lorg/w3c/dom/", "Lorg/xml/sax", "Lorg/xmlpull/v1/", "Ljunit/", "Landroidx/"};
 
        private:
            bool is_external = false;
 
            std::variant<
                EncodedMethod *,
                ExternalMethod *>
                method_encoded;
 
            mutable std::string name;
 
            mutable std::string descriptor;
 
            mutable std::string access_flag;
 
            mutable std::string class_name;
 
            mutable std::string full_name;
 
            std::uint16_t regs_from_method;
 
            std::uint16_t num_of_params;
 
            std::vector<std::unique_ptr<Instruction>> instructions;
 
            BasicBlocks basic_blocks;
 
            std::vector<Disassembler::exceptions_data> exceptions;
 
            std::vector<std::tuple<ClassAnalysis *, FieldAnalysis *, std::uint64_t>> xrefread;
            std::vector<std::tuple<ClassAnalysis *, FieldAnalysis *, std::uint64_t>> xrefwrite;
 
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> xrefto;
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> xreffrom;
 
            std::vector<std::pair<ClassAnalysis *, std::uint64_t>> xrefnewinstance;
            std::vector<std::pair<ClassAnalysis *, std::uint64_t>> xrefconstclass;
 
            void dump_instruction_dot(std::ofstream& dot_file, Instruction* instr);
 
            void dump_block_dot(std::ofstream& dot_file, DVMBasicBlock* bb);
 
            void dump_method_dot(std::ofstream& dot_file);
 
            void create_basic_blocks();
 
        public:
            MethodAnalysis(
                std::variant<EncodedMethod *, ExternalMethod *> method_encoded,
                std::vector<std::unique_ptr<Instruction>> &instructions) : method_encoded(method_encoded), instructions(std::move(instructions))
            {
                is_external = method_encoded.index() == 0 ? false : true;
 
                if (!is_external)
                {
                    auto em = std::get<EncodedMethod*>(method_encoded);
 
                    regs_from_method = em->get_code_item().get_registers_size();
 
                    num_of_params = em->getMethodID()->get_proto()->get_parameters().size();
                }
 
                if (this->instructions.size() > 0)
                    create_basic_blocks();
            }
 
            void dump_dot_file(std::string& file_path)
            {
                std::ofstream dot_file;
 
                dot_file.open(file_path);
 
                dump_method_dot(dot_file);
            }
 
            bool external() const
            {
                return is_external;
            }
 
            const BasicBlocks& get_basic_blocks() const
            {
                return basic_blocks;
            }
 
            BasicBlocks& get_basic_blocks()
            {
                return basic_blocks;
            }
 
            std::uint32_t get_number_of_registers() const
            {
                return regs_from_method;
            }
 
            std::uint32_t get_number_of_parameters() const
            {
                return num_of_params;
            }
 
            bool is_android_api() const;
 
            const std::string &get_name() const;
 
            const std::string &get_descriptor() const;
 
            const std::string &get_access_flags() const;
 
            const std::string &get_class_name() const;
 
            const std::string &get_full_name() const;
 
            std::vector<std::unique_ptr<Instruction>>& get_instructions()
            {
                return instructions;
            }
 
            std::variant<EncodedMethod *, ExternalMethod *> get_encoded_method() const
            {
                return method_encoded;
            }
 
            const std::vector<std::tuple<ClassAnalysis *, FieldAnalysis *, std::uint64_t>> &
            get_xrefread() const
            {
                return xrefread;
            }
 
            std::vector<std::tuple<ClassAnalysis *, FieldAnalysis *, std::uint64_t>> &
            get_xrefread()
            {
                return xrefread;
            }
 
            const std::vector<std::tuple<ClassAnalysis *, FieldAnalysis *, std::uint64_t>> &
            get_xrefwrite() const
            {
                return xrefwrite;
            }
 
            std::vector<std::tuple<ClassAnalysis *, FieldAnalysis *, std::uint64_t>> &
            get_xrefwrite()
            {
                return xrefwrite;
            }
 
            const std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xrefto() const
            {
                return xrefto;
            }
 
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xrefto()
            {
                return xrefto;
            }
 
            const std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xreffrom() const
            {
                return xreffrom;
            }
 
            std::vector<std::tuple<ClassAnalysis *, MethodAnalysis *, std::uint64_t>> &
            get_xreffrom()
            {
                return xreffrom;
            }
 
            const std::vector<std::pair<ClassAnalysis *, std::uint64_t>> &
            get_xrefnewinstance() const
            {
                return xrefnewinstance;
            }
 
            std::vector<std::pair<ClassAnalysis *, std::uint64_t>> &
            get_xrefnewinstance()
            {
                return xrefnewinstance;
            }
 
            const std::vector<std::pair<ClassAnalysis *, std::uint64_t>> &
            get_xrefconstclass() const
            {
                return xrefconstclass;
            }
 
            std::vector<std::pair<ClassAnalysis *, std::uint64_t>> &
            get_xrefconstclass()
            {
                return xrefconstclass;
            }
 
            void add_xrefread(ClassAnalysis *c, FieldAnalysis *f, std::uint64_t offset)
            {
                xrefread.push_back(std::make_tuple(c, f, offset));
            }
 
            void add_xrefwrite(ClassAnalysis *c, FieldAnalysis *f, std::uint64_t offset)
            {
                xrefwrite.push_back(std::make_tuple(c, f, offset));
            }
 
            void add_xrefto(ClassAnalysis *c, MethodAnalysis *m, std::uint64_t offset)
            {
                xrefto.push_back(std::make_tuple(c, m, offset));
            }
 
            void add_xreffrom(ClassAnalysis *c, MethodAnalysis *m, std::uint64_t offset)
            {
                xreffrom.push_back(std::make_tuple(c, m, offset));
            }
 
            void add_xrefnewinstance(ClassAnalysis *c, std::uint64_t offset)
            {
                xrefnewinstance.push_back(std::make_pair(c, offset));
            }
 
            void add_xrefconstclass(ClassAnalysis *c, std::uint64_t offset)
            {
                xrefconstclass.push_back(std::make_pair(c, offset));
            }
        };
 
        class ClassAnalysis
        {
        public:
            using classxref = std::unordered_map<ClassAnalysis *,
                                                 std::set<std::tuple<TYPES::REF_TYPE, MethodAnalysis *, std::uint64_t>>>;
 
        private:
            std::variant<ClassDef *, ExternalClass *> class_def;
 
            bool is_external;
 
            mutable std::string extends_;
 
            mutable std::string name_;
 
            std::unordered_map<std::string, MethodAnalysis *> methods;
            std::unordered_map<EncodedField *, std::unique_ptr<FieldAnalysis>> fields;
 
            classxref xrefto;
            classxref xreffrom;
 
            std::vector<std::pair<MethodAnalysis *, std::uint64_t>> xrefnewinstance;
 
            std::vector<std::pair<MethodAnalysis *, std::uint64_t>> xrefconstclass;
 
        public:
            ClassAnalysis(std::variant<ClassDef *, ExternalClass *> class_def) : class_def(class_def)
            {
                is_external = class_def.index() == 0 ? false : true;
            }
 
            void add_method(MethodAnalysis *method_analysis);
 
            size_t get_nb_methods() const
            {
                return methods.size();
            }
 
            size_t get_nb_fields() const
            {
                return fields.size();
            }
 
            std::variant<ClassDef *, ExternalClass *> get_class_definition() const
            {
                return class_def;
            }
 
            bool is_class_external() const
            {
                return is_external;
            }
 
            std::string& extends() const;
 
            std::string& name() const;
 
            std::vector<DVMClass *> &implements();
 
            const std::unordered_map<std::string, MethodAnalysis *> &
            get_methods() const
            {
                return methods;
            }
 
            std::unordered_map<std::string, MethodAnalysis *> &
            get_methods()
            {
                return methods;
            }
 
            const std::unordered_map<EncodedField *, std::unique_ptr<FieldAnalysis>> &
            get_fields() const
            {
                return fields;
            }
 
            std::unordered_map<EncodedField *, std::unique_ptr<FieldAnalysis>> &
            get_fields()
            {
                return fields;
            }
 
            MethodAnalysis *get_method_analysis(std::variant<EncodedMethod *, ExternalMethod *> method);
 
            FieldAnalysis *get_field_analysis(EncodedField *field);
 
            void add_field_xref_read(MethodAnalysis *method,
                                     ClassAnalysis *classobj,
                                     EncodedField *field,
                                     std::uint64_t off)
            {
                if (fields.find(field) == fields.end())
                    fields[field] = std::make_unique<FieldAnalysis>(field);
                fields[field]->add_xrefread(classobj, method, off);
            }
 
            void add_field_xref_write(MethodAnalysis *method,
                                      ClassAnalysis *classobj,
                                      EncodedField *field,
                                      std::uint64_t off)
            {
                if (fields.find(field) == fields.end())
                    fields[field] = std::make_unique<FieldAnalysis>(field);
                fields[field]->add_xrefwrite(classobj, method, off);
            }
 
            void add_method_xref_to(MethodAnalysis *method1,
                                    ClassAnalysis *classobj,
                                    MethodAnalysis *method2,
                                    std::uint64_t off)
            {
                auto method_key = method1->get_full_name();
 
                if (methods.find(method_key) == methods.end())
                    add_method(method1);
                methods[method_key]->add_xrefto(classobj, method2, off);
            }
 
            void add_method_xref_from(MethodAnalysis *method1,
                                      ClassAnalysis *classobj,
                                      MethodAnalysis *method2,
                                      std::uint64_t off)
            {
                auto method_key = method1->get_full_name();
 
                if (methods.find(method_key) == methods.end())
                    add_method(method1);
                methods[method_key]->add_xreffrom(classobj, method2, off);
            }
 
            void add_xref_to(TYPES::REF_TYPE ref_kind,
                             ClassAnalysis *classobj,
                             MethodAnalysis *methodobj,
                             std::uint64_t offset)
            {
                xrefto[classobj].insert(std::make_tuple(ref_kind, methodobj, offset));
            }
 
            void add_xref_from(TYPES::REF_TYPE ref_kind,
                               ClassAnalysis *classobj,
                               MethodAnalysis *methodobj,
                               std::uint64_t offset)
            {
                xreffrom[classobj].insert(std::make_tuple(ref_kind, methodobj, offset));
            }
 
            void add_xref_new_instance(MethodAnalysis *methodobj, std::uint64_t offset)
            {
                xrefnewinstance.push_back(std::make_pair(methodobj, offset));
            }
 
            void add_xref_const_class(MethodAnalysis *methodobj, std::uint64_t offset)
            {
                xrefconstclass.push_back(std::make_pair(methodobj, offset));
            }
 
            const classxref &
            get_xrefto() const
            {
                return xrefto;
            }
 
            classxref &
            get_xrefto()
            {
                return xrefto;
            }
 
            const classxref &
            get_xreffrom() const
            {
                return xreffrom;
            }
 
            classxref &
            get_xreffrom()
            {
                return xreffrom;
            }
 
            const std::vector<std::pair<MethodAnalysis *, std::uint64_t>>
            get_xrefnewinstance() const
            {
                return xrefnewinstance;
            }
 
            std::vector<std::pair<MethodAnalysis *, std::uint64_t>>
            get_xrefnewinstance()
            {
                return xrefnewinstance;
            }
 
            const std::vector<std::pair<MethodAnalysis *, std::uint64_t>>
            get_xrefconstclass() const
            {
                return xrefconstclass;
            }
 
            std::vector<std::pair<MethodAnalysis *, std::uint64_t>>
            get_xrefconstclass()
            {
                return xrefconstclass;
            }
        };
    } // namespace DEX
} // namespace KUNAI
 
#endif