From 87b83d754a413f690a592b7a9136ba5ba2d84953 Mon Sep 17 00:00:00 2001 From: Lil-Ran Date: Sun, 2 Mar 2025 21:39:14 +0800 Subject: [PATCH] wip: ast! :) --- ASTree.cpp | 177 ++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 176 insertions(+), 1 deletion(-) diff --git a/ASTree.cpp b/ASTree.cpp index 1f419d0..52a747b 100644 --- a/ASTree.cpp +++ b/ASTree.cpp @@ -5,6 +5,7 @@ #include "FastStack.h" #include "pyc_numeric.h" #include "bytecode.h" +#include "plusaes.hpp" // This must be a triple quote (''' or """), to handle interpolated string literals containing the opposite quote style. // E.g. f'''{"interpolated "123' literal"}''' -> valid. @@ -72,6 +73,138 @@ static void CheckIfExpr(FastStack& stack, PycRef curblock) stack.push(new ASTTernary(std::move(if_block), std::move(if_expr), std::move(else_expr))); } +PycRef PyarmorMixStrDecrypt(const std::string &inputString, PycModule *mod) +{ + std::string result(inputString.substr(1)); + if (inputString[0] & 0x80) + { + unsigned char nonce[16] = {0}; + memcpy(nonce, mod->pyarmor_mix_str_aes_nonce, 12); + nonce[15] = 2; + + plusaes::crypt_ctr( + (unsigned char *)&result[0], + result.length(), + mod->pyarmor_aes_key, + 16, + &nonce); + } + switch (inputString[0] & 0x7F) + { + case 1: + { + PycRef new_str = new PycString(PycString::TYPE_UNICODE); + new_str->setValue(result); + return new ASTObject(new_str.cast()); + } + case 2: + { + PycBuffer buf(result.data(), result.length()); + return new ASTObject(LoadObject(&buf, mod)); + } + case 3: + { + PycRef new_str = new PycString(PycString::TYPE_UNICODE); + new_str->setValue(result); + return new ASTImport(new ASTName(new_str), nullptr); + } + case 4: + default: + { + fprintf(stderr, "Unknown PyarmorAssert string first byte: %d\n", inputString[0] & 0x7F); + PycRef new_str = new PycString(PycString::TYPE_UNICODE); + new_str->setValue((char)(inputString[0] & 0x7F) + result); + return new ASTObject(new_str.cast()); + } + } +} + +void CallOrPyarmorBuiltins(FastStack &stack, PycRef &curblock, PycModule *mod) +{ + PycRef top = stack.top(); + if (top->type() != ASTNode::NODE_CALL) + return; + + PycRef call = top.cast(); + if (call->func().type() != ASTNode::NODE_OBJECT) + return; + + PycRef func_name_obj = call->func().cast()->object(); + if (func_name_obj.try_cast() == nullptr) + return; + + if (!func_name_obj.cast()->startsWith("__pyarmor_")) + return; + + const std::string& name = func_name_obj.cast()->strValue(); + if (name.find("__pyarmor_assert_") == std::string::npos) + { + PycRef new_str = new PycString(PycString::TYPE_UNICODE); + new_str->setValue(name + "(...)"); + stack.pop(); + stack.push(new ASTObject(new_str.cast())); + // str '__pyarmor_enter_12345__(...)' + return; + } + + if (call->pparams().size() != 1) // pyarmor_assert takes exactly one parameter + return; + + const auto& param = call->pparams().front(); + if (param.type() == ASTNode::NODE_OBJECT) + { + PycRef obj = param.cast()->object(); + if (obj.try_cast() == nullptr) + return; + PycRef new_node = PyarmorMixStrDecrypt(obj.cast()->strValue(), mod); + stack.pop(); + stack.push(new_node); + // result of __pyarmor_assert__(b'something') + return; + } + + if (param.type() == ASTNode::NODE_TUPLE) + { + const auto &tuple = param.cast(); + if (tuple->values().size() <= 1 || tuple->values().size() > 3) + return; + if (tuple->values()[1].type() != ASTNode::NODE_OBJECT) + return; + auto enc_str = tuple->values()[1].cast()->object().try_cast(); + if (enc_str == nullptr) + return; + PycRef attr_name = PyarmorMixStrDecrypt(enc_str->strValue(), mod); + if (attr_name->type() != ASTNode::NODE_OBJECT) + return; + auto name_str = attr_name.cast()->object().try_cast(); + if (name_str == nullptr) + return; + PycRef attr_ref = new ASTBinary(tuple->values()[0], new ASTName(name_str), ASTBinary::BIN_ATTR); + if (tuple->values().size() == 2) + { + stack.pop(); + stack.push(attr_ref); + // __pyarmor_assert__((from, b'enc_attr')) -> from.enc_attr + return; + } + else if (tuple->values().size() == 3) + { + stack.pop(); + curblock->append(new ASTStore(tuple->values()[2], attr_ref)); + // __pyarmor_assert__((from, b'enc_attr', value)) -> from.enc_attr = value + return; + } + } + + if (param.type() == ASTNode::NODE_NAME) + { + stack.pop(); + stack.push(param); + // __pyarmor_assert__(name) -> name + return; + } +} + PycRef BuildFromCode(PycRef code, PycModule* mod) { PycBuffer source(code->code()->value(), code->code()->length()); @@ -521,6 +654,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) } stack.push(new ASTCall(func, pparamList, kwparamList)); + + CallOrPyarmorBuiltins(stack, curblock, mod); } break; case Pyc::CALL_FUNCTION_VAR_A: @@ -548,6 +683,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) PycRef call = new ASTCall(func, pparamList, kwparamList); call.cast()->setVar(var); stack.push(call); + + CallOrPyarmorBuiltins(stack, curblock, mod); } break; case Pyc::CALL_FUNCTION_KW_A: @@ -575,6 +712,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) PycRef call = new ASTCall(func, pparamList, kwparamList); call.cast()->setKW(kw); stack.push(call); + + CallOrPyarmorBuiltins(stack, curblock, mod); } break; case Pyc::CALL_FUNCTION_VAR_KW_A: @@ -605,6 +744,30 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) call.cast()->setKW(kw); call.cast()->setVar(var); stack.push(call); + + CallOrPyarmorBuiltins(stack, curblock, mod); + } + break; + case Pyc::CALL_FUNCTION_EX_A: + { + PycRef kw; + if (operand & 0x01) + { + kw = stack.top(); + stack.pop(); + } + PycRef var = stack.top(); + stack.pop(); + PycRef func = stack.top(); + stack.pop(); + + PycRef call = new ASTCall(func, ASTCall::pparam_t(), ASTCall::kwparam_t()); + if (operand & 0x01) + call.cast()->setKW(kw); + call.cast()->setVar(var); + stack.push(call); + + CallOrPyarmorBuiltins(stack, curblock, mod); } break; case Pyc::CALL_METHOD_A: @@ -633,6 +796,8 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) PycRef func = stack.top(); stack.pop(); stack.push(new ASTCall(func, pparamList, ASTCall::kwparam_t())); + + CallOrPyarmorBuiltins(stack, curblock, mod); } break; case Pyc::CONTINUE_LOOP_A: @@ -1307,6 +1472,16 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) bool push = true; do { + // A polyfill for Pyarmor which jumps forward on top level scope + auto &top = blocks.top(); + if (top == defblock) { + pos += offs; + for (int i = 0; i < offs; i++) { + source.getByte(); + } + break; + } + blocks.pop(); if (!blocks.empty()) @@ -2474,7 +2649,7 @@ PycRef BuildFromCode(PycRef code, PycModule* mod) } break; default: - fprintf(stderr, "Unsupported opcode: %s (%d)\n", Pyc::OpcodeName(opcode), opcode); + fprintf(stderr, "Unsupported opcode: %s (%d) at %s\n", Pyc::OpcodeName(opcode), opcode, code->qualName()->value()); cleanBuild = false; return new ASTNodeList(defblock->nodes()); }