python-assembler 0.1.1

Parse and write Python .pyc bytecode
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215

//! # Marshal 序列化模块
//!
//! 本模块提供了对 Python 对象进行 Marshal 序列化的功能，主要用于生成 `.pyc` 文件。
//! 它实现了将 `PythonObject` 和 `PythonCodeObject` 转换为字节流的逻辑。

use crate::{
    instructions::PythonInstruction,
    program::{PythonCodeObject, PythonObject, PythonVersion},
};
use gaia_binary::{BinaryWriter, Fixed, LittleEndian};
use std::io::{self, Write};

/// Marshal 序列化器
#[derive(Debug)]
pub struct MarshalWriter<W: Write> {
    writer: BinaryWriter<W, Fixed<LittleEndian>>,
    #[allow(dead_code)]
    version: PythonVersion,
}

impl<W: Write> MarshalWriter<W> {
    /// 创建一个新的 MarshalWriter 实例
    pub fn new(writer: W, version: PythonVersion) -> Self {
        Self { writer: BinaryWriter::new(writer), version }
    }

    /// 写入一个 Python 对象
    pub fn write_object(&mut self, obj: &PythonObject) -> io::Result<()> {
        match obj {
            PythonObject::None => {
                self.writer.write_u8(b'N').map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
            }
            PythonObject::Bool(b) => {
                self.writer
                    .write_u8(if *b { b'T' } else { b'F' })
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
            }
            PythonObject::Int(i) => {
                self.writer.write_u8(b'i' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_INT | FLAG_REF
                self.writer.write_i32(*i).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
            }
            PythonObject::Integer(i) => {
                if *i >= i32::MIN as i64 && *i <= i32::MAX as i64 {
                    self.writer.write_u8(b'i' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                    self.writer.write_i32(*i as i32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                }
                else {
                    self.writer.write_u8(b'I' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_INT64 | FLAG_REF
                    self.writer.write_i64(*i).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                }
            }
            PythonObject::Float(f) => {
                self.writer.write_u8(b'g' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_BINARY_FLOAT | FLAG_REF
                self.writer.write_f64(*f).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
            }
            PythonObject::Str(s) | PythonObject::String(s) => {
                if s.len() <= 255 && s.is_ascii() {
                    self.writer.write_u8(b'Z' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_SHORT_ASCII_INTERNED | FLAG_REF
                    self.writer.write_u8(s.len() as u8).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                    self.writer.write_bytes(s.as_bytes()).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                }
                else {
                    self.writer.write_u8(b'u' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_UNICODE | FLAG_REF
                    self.writer.write_u32(s.len() as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                    self.writer.write_bytes(s.as_bytes()).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                }
            }
            PythonObject::Bytes(b) => {
                self.writer.write_u8(b's' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_STRING | FLAG_REF
                self.writer.write_u32(b.len() as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                self.writer.write_bytes(b).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
            }
            PythonObject::List(l) => {
                self.writer.write_u8(b'[' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_LIST | FLAG_REF
                self.writer.write_u32(l.len() as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                for item in l {
                    self.write_object(item)?;
                }
            }
            PythonObject::Tuple(t) => {
                if t.len() <= 255 {
                    self.writer.write_u8(b')' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_SMALL_TUPLE | FLAG_REF
                    self.writer.write_u8(t.len() as u8).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                }
                else {
                    self.writer.write_u8(b'(' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_TUPLE | FLAG_REF
                    self.writer.write_u32(t.len() as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                }
                for item in t {
                    self.write_object(item)?;
                }
            }
            PythonObject::Code(code) => {
                self.write_code_object(code)?;
            }
        }
        Ok(())
    }

    /// 写入一个 Python 代码对象
    pub fn write_code_object(&mut self, code: &PythonCodeObject) -> io::Result<()> {
        self.writer.write_u8(b'c' | 0x80).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; // TYPE_CODE | FLAG_REF

        // Python 3.11+ code object fields
        self.writer.write_u32(code.co_argcount as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        self.writer.write_u32(code.co_posonlyargcount as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        self.writer.write_u32(code.co_kwonlyargcount as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        // nlocals removed in 3.11+ marshal, replaced by stacksize and flags
        self.writer.write_u32(code.co_stacksize as u32).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        self.writer.write_u32(code.co_flags).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        // co_code
        let bytecode = self.encode_instructions(&code.co_code);
        self.write_object(&PythonObject::Bytes(bytecode))?;

        // co_consts
        self.write_object(&PythonObject::Tuple(code.co_consts.clone()))?;

        // co_names
        self.write_object(&PythonObject::Tuple(code.co_names.iter().map(|s| PythonObject::Str(s.clone())).collect()))?;

        // co_localsplusnames (Python 3.11+)
        self.write_object(&PythonObject::Tuple(
            code.co_localsplusnames.iter().map(|s| PythonObject::Str(s.clone())).collect(),
        ))?;

        // co_localspluskinds
        self.write_object(&PythonObject::Bytes(code.co_localspluskinds.clone()))?;

        // co_filename
        self.write_object(&PythonObject::Str(code.source_name.clone()))?;

        // co_name
        self.write_object(&PythonObject::Str(code.name.clone()))?;

        // co_qualname
        self.write_object(&PythonObject::Str(code.qualname.clone()))?;

        // co_firstlineno
        self.writer.write_u32(code.first_line).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        // co_linetable
        self.write_object(&PythonObject::Bytes(code.co_linetable.clone()))?;

        // co_exceptiontable
        self.write_object(&PythonObject::Bytes(code.co_exceptiontable.clone()))?;

        Ok(())
    }

    fn encode_instructions(&self, instrs: &[PythonInstruction]) -> Vec<u8> {
        let mut data = Vec::new();
        for instr in instrs {
            let (opcode, arg) = self.map_instruction(instr);
            data.push(opcode);
            data.push((arg & 0xFF) as u8);
        }
        data
    }

    fn map_instruction(&self, instr: &PythonInstruction) -> (u8, u32) {
        match instr {
            PythonInstruction::RESUME => (151, 0),
            PythonInstruction::RETURN_CONST(idx) => (121, *idx),
            PythonInstruction::LOAD_CONST(idx) => (100, *idx),
            PythonInstruction::STORE_NAME(idx) => (90, *idx),
            PythonInstruction::LOAD_NAME(idx) => (101, *idx),
            PythonInstruction::PUSH_NULL => (2, 0),
            PythonInstruction::CALL(argc) => (171, *argc),
            PythonInstruction::BINARY_OP(op) => (122, *op),
            PythonInstruction::POP_TOP => (1, 0),
            PythonInstruction::RETURN_VALUE => (83, 0),
            PythonInstruction::BINARY_MODULO => (22, 0),
            PythonInstruction::INPLACE_MODULO => (59, 0),
            PythonInstruction::NOP => (9, 0),
            PythonInstruction::UNARY_NEGATIVE => (11, 0),
            PythonInstruction::UNARY_NOT => (12, 0),
            PythonInstruction::UNARY_INVERT => (15, 0),
            PythonInstruction::GET_LEN => (30, 0),
            PythonInstruction::MATCH_MAPPING => (31, 0),
            PythonInstruction::MATCH_SEQUENCE => (32, 0),
            PythonInstruction::MATCH_KEYS => (33, 0),
            PythonInstruction::GET_ITER => (68, 0),
            PythonInstruction::LOAD_BUILD_CLASS => (71, 0),
            PythonInstruction::LOAD_ASSERTION_ERROR => (74, 0),
            PythonInstruction::RETURN_GENERATOR => (75, 0),
            PythonInstruction::YIELD_VALUE => (150, 0),
            PythonInstruction::LOAD_FAST(idx) => (124, *idx),
            PythonInstruction::STORE_FAST(idx) => (125, *idx),
            PythonInstruction::DELETE_FAST(idx) => (126, *idx),
            PythonInstruction::LOAD_GLOBAL(idx) => (116, *idx),
            PythonInstruction::STORE_GLOBAL(idx) => (97, *idx),
            PythonInstruction::DELETE_GLOBAL(idx) => (98, *idx),
            PythonInstruction::LOAD_ATTR(idx) => (106, *idx),
            PythonInstruction::STORE_ATTR(idx) => (95, *idx),
            PythonInstruction::DELETE_ATTR(idx) => (96, *idx),
            PythonInstruction::COMPARE_OP(idx) => (107, *idx),
            PythonInstruction::IMPORT_NAME(idx) => (108, *idx),
            PythonInstruction::IMPORT_FROM(idx) => (109, *idx),
            PythonInstruction::JUMP_FORWARD(idx) => (110, *idx),
            PythonInstruction::JUMP_BACKWARD(idx) => (140, *idx),
            PythonInstruction::POP_JUMP_IF_FALSE(idx) => (114, *idx),
            PythonInstruction::POP_JUMP_IF_TRUE(idx) => (115, *idx),
            PythonInstruction::MAKE_FUNCTION(idx) => (132, *idx),
            PythonInstruction::BUILD_TUPLE(idx) => (102, *idx),
            PythonInstruction::BUILD_LIST(idx) => (103, *idx),
            PythonInstruction::BUILD_SET(idx) => (104, *idx),
            PythonInstruction::BUILD_MAP(idx) => (105, *idx),
            PythonInstruction::LIST_APPEND(idx) => (145, *idx),
            PythonInstruction::SET_ADD(idx) => (146, *idx),
            PythonInstruction::MAP_ADD(idx) => (147, *idx),
            _ => (0, 0), // Default to NOP or handle error
        }
    }
}