/* * Python struct module. * * renamed to pystruct since struct is a rust keyword. * * Use this rust module to do byte packing: * https://docs.rs/byteorder/1.2.6/byteorder/ */ extern crate byteorder; use self::byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; use super::super::obj::{objbool, objbytes, objfloat, objint, objstr, objtype}; use super::super::pyobject::{ DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol, }; use super::super::VirtualMachine; use num_bigint::{BigInt, ToBigInt}; use num_traits::ToPrimitive; use std::io::{Cursor, Read, Write}; #[derive(Debug)] struct FormatCode { code: char, repeat: i32, } fn parse_format_string(fmt: String) -> Vec { // First determine "<", ">","!" or "=" // TODO // Now, analyze struct string furter: let mut codes = vec![]; for c in fmt.chars() { match c { 'b' | 'B' | 'h' | 'H' | 'i' | 'I' | 'q' | 'Q' | 'f' | 'd' => { codes.push(FormatCode { code: c, repeat: 1 }) } c => { panic!("Illegal format code {:?}", c); } } } codes } fn get_int(vm: &mut VirtualMachine, arg: &PyObjectRef) -> Result { objint::to_int(vm, arg, 10) } fn pack_i8( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_i8().unwrap(); data.write_i8(v).unwrap(); Ok(()) } fn pack_u8( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_u8().unwrap(); data.write_u8(v).unwrap(); Ok(()) } fn pack_bool( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { if objtype::isinstance(&arg, &vm.ctx.bool_type()) { let v = if objbool::get_value(arg) { 1 } else { 0 }; data.write_u8(v).unwrap(); Ok(()) } else { Err(vm.new_type_error(format!("Expected boolean"))) } } fn pack_i16( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_i16().unwrap(); data.write_i16::(v).unwrap(); Ok(()) } fn pack_u16( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_u16().unwrap(); data.write_u16::(v).unwrap(); Ok(()) } fn pack_i32( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_i32().unwrap(); data.write_i32::(v).unwrap(); Ok(()) } fn pack_u32( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_u32().unwrap(); data.write_u32::(v).unwrap(); Ok(()) } fn pack_i64( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_i64().unwrap(); data.write_i64::(v).unwrap(); Ok(()) } fn pack_u64( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { let v = get_int(vm, arg)?.to_u64().unwrap(); data.write_u64::(v).unwrap(); Ok(()) } fn pack_f32( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { if objtype::isinstance(&arg, &vm.ctx.float_type()) { let v = objfloat::get_value(arg) as f32; data.write_f32::(v).unwrap(); Ok(()) } else { Err(vm.new_type_error(format!("Expected float"))) } } fn pack_f64( vm: &mut VirtualMachine, arg: &PyObjectRef, data: &mut Write, ) -> Result<(), PyObjectRef> { if objtype::isinstance(&arg, &vm.ctx.float_type()) { let v = objfloat::get_value(arg) as f64; data.write_f64::(v).unwrap(); Ok(()) } else { Err(vm.new_type_error(format!("Expected float"))) } } fn struct_pack(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { if args.args.len() < 1 { Err(vm.new_type_error(format!( "Expected at least 1 argument (got: {})", args.args.len() ))) } else { let fmt_arg = args.args[0].clone(); if objtype::isinstance(&fmt_arg, &vm.ctx.str_type()) { let fmt_str = objstr::get_value(&fmt_arg); let codes = parse_format_string(fmt_str); if codes.len() + 1 == args.args.len() { // Create data vector: let mut data = Vec::::new(); // Loop over all opcodes: for (code, arg) in codes.iter().zip(args.args.iter().skip(1)) { debug!("code: {:?}", code); match code.code { 'b' => pack_i8(vm, arg, &mut data)?, 'B' => pack_u8(vm, arg, &mut data)?, '?' => pack_bool(vm, arg, &mut data)?, 'h' => pack_i16(vm, arg, &mut data)?, 'H' => pack_u16(vm, arg, &mut data)?, 'i' => pack_i32(vm, arg, &mut data)?, 'I' => pack_u32(vm, arg, &mut data)?, 'l' => pack_i32(vm, arg, &mut data)?, 'L' => pack_u32(vm, arg, &mut data)?, 'q' => pack_i64(vm, arg, &mut data)?, 'Q' => pack_u64(vm, arg, &mut data)?, 'f' => pack_f32(vm, arg, &mut data)?, 'd' => pack_f64(vm, arg, &mut data)?, c => { panic!("Unsupported format code {:?}", c); } } } Ok(vm.ctx.new_bytes(data)) } else { Err(vm.new_type_error(format!( "Expected {} arguments (got: {})", codes.len() + 1, args.args.len() ))) } } else { Err(vm.new_type_error(format!("First argument must be of str type"))) } } } fn unpack_i8(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_i8() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_u8(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_u8() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_bool(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_u8() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_bool(v > 0)), } } fn unpack_i16(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_i16::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_u16(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_u16::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_i32(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_i32::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_u32(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_u32::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_i64(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_i64::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_u64(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_u64::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_int(v.to_bigint().unwrap())), } } fn unpack_f32(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_f32::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_float(v as f64)), } } fn unpack_f64(vm: &mut VirtualMachine, rdr: &mut Read) -> PyResult { match rdr.read_f64::() { Err(err) => panic!("Error in reading {:?}", err), Ok(v) => Ok(vm.ctx.new_float(v)), } } fn struct_unpack(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, args, required = [ (fmt, Some(vm.ctx.str_type())), (buffer, Some(vm.ctx.bytes_type())) ] ); let fmt_str = objstr::get_value(&fmt); let codes = parse_format_string(fmt_str); let data = objbytes::get_value(buffer).to_vec(); let mut rdr = Cursor::new(data); let mut items = vec![]; for code in codes { debug!("unpack code: {:?}", code); match code.code { 'b' => items.push(unpack_i8(vm, &mut rdr)?), 'B' => items.push(unpack_u8(vm, &mut rdr)?), '?' => items.push(unpack_bool(vm, &mut rdr)?), 'h' => items.push(unpack_i16(vm, &mut rdr)?), 'H' => items.push(unpack_u16(vm, &mut rdr)?), 'i' => items.push(unpack_i32(vm, &mut rdr)?), 'I' => items.push(unpack_u32(vm, &mut rdr)?), 'l' => items.push(unpack_i32(vm, &mut rdr)?), 'L' => items.push(unpack_u32(vm, &mut rdr)?), 'q' => items.push(unpack_i64(vm, &mut rdr)?), 'Q' => items.push(unpack_u64(vm, &mut rdr)?), 'f' => items.push(unpack_f32(vm, &mut rdr)?), 'd' => items.push(unpack_f64(vm, &mut rdr)?), c => { panic!("Unsupported format code {:?}", c); } } } Ok(vm.ctx.new_tuple(items)) } pub fn mk_module(ctx: &PyContext) -> PyObjectRef { let py_mod = ctx.new_module(&"struct".to_string(), ctx.new_scope(None)); py_mod.set_item("pack", ctx.new_rustfunc(struct_pack)); py_mod.set_item("unpack", ctx.new_rustfunc(struct_unpack)); py_mod }