sile / amf Goto Github PK

// amf_packet.rs

use std::{io};
use std::convert::{TryFrom, TryInto};
use amf::{Version, Amf0Value};
use amf::amf0::{Value, array, string, number};
extern crate byteorder;
use byteorder::{ReadBytesExt, WriteBytesExt, BigEndian, LittleEndian};
use std::fs::File;
use std::io::Read;
use std::string::FromUtf8Error;
use std::borrow::BorrowMut;

pub struct AmfHeader {
    pub name: String,
    pub must_understand: bool,
    pub content: Value,
}

pub struct AmfMessage {
    pub target: String,
    pub response: String,
    pub content: Value,
}

pub struct AmfPacket {
    pub version: Version,
    pub headers: Vec<AmfHeader>,
    pub messages: Vec<AmfMessage>,
}

trait ReadBytesExt2: io::Read {
    #[inline]
    fn read_bytes(&mut self, size: usize) -> Vec<u8> {
        let mut buf = vec![0; size];
        self.read_exact(&mut buf).unwrap();
        buf
    }

    #[inline]
    fn read_utf8_string(&mut self, length: usize) -> io::Result<String> {
        Ok(String::from_utf8(self.read_bytes(length)).unwrap())
    }
}
impl<R: io::Read + ?Sized> ReadBytesExt2 for R {}

impl AmfPacket {
    pub fn new(version: Version) -> Self {
        AmfPacket {version, headers: Vec::new(), messages: Vec::new()}
    }

    pub fn read_from<R>(mut reader: R) -> io::Result<Self>
        where
            R: io::Read,
    {
        let version = reader.read_u16::<BigEndian>()?;
        let mut headers = Vec::new();
        let header_count = reader.read_u16::<BigEndian>()?;
        for _ in 0..header_count {
            let header_name_length = reader.read_u16::<BigEndian>()?;
            let header_name = reader.read_utf8_string(header_name_length as usize)?;
            let must_understand = reader.read_u8()? != 0;
            let _header_length = reader.read_i32::<BigEndian>()?;
            let content = Value::read_from(&mut reader).unwrap();
            headers.push(AmfHeader {name: header_name, must_understand, content });
        }

        let mut messages = Vec::new();
        let message_count = reader.read_u16::<BigEndian>()?;
        for _ in 0..message_count {
            let target_length = reader.read_u16::<BigEndian>()?;
            let target = reader.read_utf8_string(target_length as usize)?;
            let response_length = reader.read_u16::<BigEndian>()?;
            let response = reader.read_utf8_string(response_length as usize)?;
            let _message_length = reader.read_i32::<BigEndian>()?;
            let content = Value::read_from(&mut reader).unwrap();
            messages.push(AmfMessage {target, response, content});
        }

        Ok(AmfPacket { version: match version {
            0 => Version::Amf0,
            3 => Version::Amf3,
            _ => panic!("Unknown version! Probably not AMF packet: {}", version),
        }, headers, messages })
    }

    pub fn write_to<W>(&self, mut writer: W) -> io::Result<()>
        where
            W: io::Write,
    {
        writer.write_u16::<BigEndian>(match self.version {
            Version::Amf0 => 0,
            Version::Amf3 => 3
        })?;
        writer.write_u16::<BigEndian>(self.headers.len() as u16)?;

        for header in &self.headers {
            writer.write_u16::<BigEndian>(header.name.len() as u16)?;
            writer.write_all(header.name.as_bytes())?;
            writer.write_u8(header.must_understand as u8)?;
            writer.write_i32::<BigEndian>(-1)?; // https://www.adobe.com/content/dam/acom/en/devnet/pdf/amf0-file-format-specification.pdf -1 can be specified
            header.content.write_to(&mut writer)?;
        }

        writer.write_u16::<BigEndian>(self.messages.len() as u16)?;

        for message in &self.messages {
            writer.write_u16::<BigEndian>(message.target.len() as u16)?;
            writer.write_all(message.target.as_bytes())?;
            writer.write_u16::<BigEndian>(message.response.len() as u16)?;
            writer.write_all(message.response.as_bytes())?;
            writer.write_i32::<BigEndian>(-1)?;
            message.content.write_to(&mut writer)?;
        }

        Ok(())
    }
}

    /// Decodes a AMF3 value. Keeps traits/strings/complexes.
    pub fn decode_keep_state(&mut self) -> DecodeResult<Value> {
        self.decode_value()
    }