playready/
pssh.rs

1//! PSSH module.
2
3use crate::{
4    binary_format::pssh::{PSSHBox, PlayreadyHeader},
5    xml_utils,
6};
7use base64::prelude::*;
8use binrw::BinRead;
9use std::io::Cursor;
10
11/// Version of Wrm header
12pub type WrmHeaderVersion = [u8; 4];
13
14#[derive(Debug, Clone)]
15/// Wrm header which is extracted from PSSH box.
16pub struct WrmHeader {
17    header: String,
18    version: WrmHeaderVersion,
19}
20
21impl WrmHeader {
22    /// Returns version of Wrm header
23    pub fn version(&self) -> WrmHeaderVersion {
24        self.version
25    }
26
27    fn parse_version(version: String) -> Result<WrmHeaderVersion, crate::Error> {
28        if !version.chars().all(|c| c.is_ascii_digit() || c == '.') {
29            return Err(crate::Error::InvalidWrmHeader(
30                "Unexpected character in version",
31                version.into(),
32            ));
33        }
34
35        let mut result = WrmHeaderVersion::default();
36        let mut split = version.split('.');
37
38        for n in &mut result {
39            let Some(part) = split.next() else {
40                return Err(crate::Error::InvalidWrmHeader(
41                    "Not enough parts in version",
42                    version.into(),
43                ));
44            };
45
46            match part.parse() {
47                Ok(i) => *n = i,
48                Err(_) => {
49                    return Err(crate::Error::InvalidWrmHeader(
50                        "Failed to parse version number",
51                        version.into(),
52                    ));
53                }
54            }
55        }
56
57        if split.next().is_some() {
58            return Err(crate::Error::InvalidWrmHeader(
59                "Too many parts in version",
60                version.into(),
61            ));
62        }
63
64        Ok(result)
65    }
66}
67
68impl TryFrom<String> for WrmHeader {
69    type Error = crate::Error;
70
71    fn try_from(value: String) -> Result<Self, Self::Error> {
72        let Some((version, ..)) = xml_utils::parse_wrm_header(&value)? else {
73            return Err(Self::Error::InvalidWrmHeader(
74                "Failed to parse XML",
75                value.into(),
76            ));
77        };
78
79        let version = Self::parse_version(version)?;
80
81        Ok(WrmHeader {
82            header: value,
83            version,
84        })
85    }
86}
87
88impl From<WrmHeader> for String {
89    fn from(value: WrmHeader) -> Self {
90        value.header
91    }
92}
93
94/// Wrapper for `PlayreadyObject` binary format.
95#[derive(Debug, Clone)]
96pub struct Pssh {
97    parsed: PlayreadyHeader,
98}
99
100impl Pssh {
101    /// Creates [`Pssh`] from bytes.
102    pub fn from_bytes(b: &[u8]) -> Result<Self, binrw::Error> {
103        let pssh_box = PSSHBox::read(&mut Cursor::new(b));
104
105        match pssh_box {
106            Ok(pssh_box) => Ok(Self {
107                parsed: pssh_box.data,
108            }),
109            Err(_) => Ok(Self {
110                parsed: PlayreadyHeader::read(&mut Cursor::new(b))?,
111            }),
112        }
113    }
114
115    /// Creates [`Pssh`] from Base64 encoded bytes.
116    pub fn from_b64(b64: &[u8]) -> Result<Self, crate::Error> {
117        let bytes = BASE64_STANDARD.decode(b64)?;
118        Self::from_bytes(&bytes).map_err(|e| e.into())
119    }
120
121    /// Returns WRM headers parsed from `PSSHBox` or `PlayreadyObject` format.
122    pub fn wrm_headers(&self) -> Vec<WrmHeader> {
123        self.parsed
124            .records
125            .iter()
126            .filter(|o| o.type_ == 1)
127            .filter_map(|o| {
128                String::from_utf16(&o.data)
129                    .inspect_err(|e| {
130                        log::error!("Failed to create uf16 string from wrm header: {e:?}")
131                    })
132                    .ok()
133                    .and_then(|h| {
134                        WrmHeader::try_from(h)
135                            .inspect_err(|e| {
136                                log::error!("Failed to create wrm header from string: {e:?}")
137                            })
138                            .ok()
139                    })
140            })
141            .collect()
142    }
143}
144
145impl TryFrom<&[u8]> for Pssh {
146    type Error = binrw::Error;
147
148    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
149        Self::from_bytes(value)
150    }
151}
152
153#[cfg(test)]
154mod test {
155    use crate::pssh::WrmHeader;
156
157    #[test]
158    fn parse_wrmheader_with_valid_version() {
159        assert_eq!(
160            WrmHeader::try_from(String::from("<WRMHEADER version=\"1.2.3.4\" />"))
161                .unwrap()
162                .version(),
163            [1, 2, 3, 4]
164        );
165    }
166
167    #[test]
168    fn parse_wrmheader_with_invalid_xml() {
169        let header = WrmHeader::try_from(String::from("<invalid"));
170
171        assert!(matches!(header, Err(crate::Error::XmlParserError(_))));
172    }
173
174    #[test]
175    fn parse_wrmheader_with_xml_without_tag() {
176        let header = WrmHeader::try_from(String::from("<NOT_WRMHEADER />"));
177
178        assert!(matches!(header, Err(crate::Error::InvalidWrmHeader(_, _))));
179    }
180
181    #[test]
182    fn parse_wrmheader_with_invalid_version_1() {
183        let header = WrmHeader::try_from(String::from("<WRMHEADER version=\"1.2.3.4.\" />"));
184
185        assert!(matches!(header, Err(crate::Error::InvalidWrmHeader(_, _))));
186    }
187
188    #[test]
189    fn parse_wrmheader_with_invalid_version_2() {
190        let header = WrmHeader::try_from(String::from("<WRMHEADER version=\"+1.2.3.4\" />"));
191
192        assert!(matches!(header, Err(crate::Error::InvalidWrmHeader(_, _))));
193    }
194}