1use crate::{
4 binary_format::pssh::{PSSHBox, PlayreadyHeader},
5 xml_utils,
6};
7use base64::prelude::*;
8use binrw::BinRead;
9use std::io::Cursor;
10
11pub type WrmHeaderVersion = [u8; 4];
13
14#[derive(Debug, Clone)]
15pub struct WrmHeader {
17 header: String,
18 version: WrmHeaderVersion,
19}
20
21impl WrmHeader {
22 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#[derive(Debug, Clone)]
96pub struct Pssh {
97 parsed: PlayreadyHeader,
98}
99
100impl Pssh {
101 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 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 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}