285 lines
8.5 KiB
Rust
285 lines
8.5 KiB
Rust
use indicatif::ProgressBar;
|
|
use indicatif::ProgressBarIter;
|
|
use png;
|
|
use pyo3::exceptions::PyValueError;
|
|
use pyo3::prelude::*;
|
|
use std::fs::{File, OpenOptions};
|
|
use std::io::BufWriter;
|
|
use std::path::Path;
|
|
|
|
use pcap_parser::traits::PcapReaderIterator;
|
|
use pcap_parser::*;
|
|
|
|
struct PacketsIterator {
|
|
cap: PcapNGReader<ProgressBarIter<File>>,
|
|
i: usize,
|
|
size: usize,
|
|
}
|
|
|
|
impl Iterator for PacketsIterator {
|
|
type Item = Vec<u8>;
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
let mut ret = None;
|
|
while ret.is_none() {
|
|
match self.cap.next() {
|
|
Ok((offset, packet)) => {
|
|
self.i += 1;
|
|
self.size += offset;
|
|
match packet {
|
|
PcapBlockOwned::Legacy(_block) => {
|
|
println!("dunno");
|
|
}
|
|
PcapBlockOwned::LegacyHeader(_block) => {
|
|
println!("dunnoheader");
|
|
}
|
|
PcapBlockOwned::NG(block) => {
|
|
if block.is_data_block() {
|
|
match block {
|
|
Block::EnhancedPacket(ref epb) => {
|
|
if epb.origlen == 6972 {
|
|
// remove udp header
|
|
ret = Some(epb.data[0x2a..].to_vec());
|
|
}
|
|
}
|
|
Block::SimplePacket(ref ep) => {
|
|
if ep.origlen == 6972 {
|
|
println!("found one regular");
|
|
}
|
|
}
|
|
_ => {
|
|
println!("unsupported packet");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
self.cap.consume(offset)
|
|
}
|
|
Err(PcapError::Eof) => break,
|
|
Err(PcapError::Incomplete) => {
|
|
self.cap.refill().unwrap();
|
|
}
|
|
Err(_) => {
|
|
println!("unexpected error");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
ret
|
|
}
|
|
}
|
|
|
|
impl PacketsIterator {
|
|
fn new(filename: &str) -> anyhow::Result<Self> {
|
|
let file = OpenOptions::new()
|
|
.read(true)
|
|
.write(false)
|
|
.create(false)
|
|
.open(filename)?;
|
|
let pb = ProgressBar::new(file.metadata()?.len());
|
|
let wrap = pb.wrap_read(file);
|
|
let cap = PcapNGReader::new(65535, wrap)?;
|
|
Ok(PacketsIterator { cap, i: 0, size: 0 })
|
|
}
|
|
}
|
|
|
|
#[repr(C, packed)]
|
|
#[derive(Clone, Debug)]
|
|
pub struct Header {
|
|
c1: u32,
|
|
c2: u16,
|
|
part: u16,
|
|
a: u16,
|
|
ffaa: u16,
|
|
b: u16,
|
|
c: u16,
|
|
d: u16,
|
|
}
|
|
|
|
impl Header {
|
|
fn read(data: &[u8]) -> anyhow::Result<Self> {
|
|
Ok(Header {
|
|
c1: u32::from_be_bytes([data[0], data[1], data[2], data[3]]),
|
|
c2: u16::from_be_bytes([data[4], data[5]]),
|
|
part: u16::from_be_bytes([data[6], data[7]]),
|
|
a: u16::from_be_bytes([data[8], data[9]]),
|
|
ffaa: u16::from_be_bytes([data[10], data[11]]),
|
|
b: u16::from_be_bytes([data[12], data[13]]),
|
|
c: u16::from_be_bytes([data[14], data[15]]),
|
|
d: u16::from_be_bytes([data[16], data[17]]),
|
|
})
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
fn read_via_cast(data: &[u8]) -> anyhow::Result<&Self> {
|
|
let size = std::mem::size_of::<Self>();
|
|
if data.len() < size {
|
|
return Err(anyhow::anyhow!("not large enough"));
|
|
}
|
|
let mem = &data[..size].as_ptr();
|
|
Ok(unsafe { &*mem.cast() })
|
|
}
|
|
}
|
|
|
|
const HDR_SIZE: usize = std::mem::size_of::<Header>();
|
|
|
|
pub struct Frame {
|
|
#[allow(dead_code)]
|
|
header: Header,
|
|
raw: Vec<u8>,
|
|
}
|
|
|
|
impl Frame {
|
|
fn pixels(&self) -> Vec<u16> {
|
|
(0..self.raw.len())
|
|
.step_by(2)
|
|
.map(|i| u16::from_be_bytes([self.raw[i], self.raw[i + 1]]))
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
pub struct Decoder {
|
|
packiter: PacketsIterator,
|
|
parts: Vec<Vec<u8>>,
|
|
}
|
|
|
|
impl Decoder {
|
|
pub fn new(filename: &str) -> anyhow::Result<Self> {
|
|
let packiter = PacketsIterator::new(filename)?;
|
|
Ok(Decoder {
|
|
packiter,
|
|
parts: vec![],
|
|
})
|
|
}
|
|
}
|
|
|
|
impl Iterator for Decoder {
|
|
type Item = Frame;
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
let mut ret = None;
|
|
while let Some(packet) = self.packiter.next() {
|
|
let hdr = match Header::read(&packet) {
|
|
Ok(header) => header,
|
|
Err(e) => {
|
|
println!("error reading header: {:?}", e);
|
|
return None;
|
|
}
|
|
};
|
|
let data = packet[HDR_SIZE..].to_vec();
|
|
if hdr.part == 0 && self.parts.len() > 0 {
|
|
ret = Some(Frame {
|
|
header: hdr,
|
|
raw: self.parts.concat(),
|
|
});
|
|
self.parts.clear();
|
|
}
|
|
self.parts.push(data);
|
|
if ret.is_some() {
|
|
return ret;
|
|
}
|
|
}
|
|
None
|
|
}
|
|
}
|
|
|
|
fn write_raw_frame(name: &str, data: &[u8]) -> anyhow::Result<()> {
|
|
let path = Path::new(&name);
|
|
let file = File::create(path)?;
|
|
let ref mut w = BufWriter::new(file);
|
|
let mut encoder = png::Encoder::new(w, 288, 384);
|
|
encoder.set_color(png::ColorType::Grayscale);
|
|
encoder.set_depth(png::BitDepth::Sixteen);
|
|
let mut writer = encoder.write_header()?;
|
|
writer.write_image_data(data)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn write_calibrated_frame(name: &str, data: &[u16]) -> anyhow::Result<()> {
|
|
let path = Path::new(&name);
|
|
let file = File::create(path).unwrap();
|
|
let ref mut w = BufWriter::new(file);
|
|
let mut encoder = png::Encoder::new(w, 288, 384);
|
|
encoder.set_color(png::ColorType::Grayscale);
|
|
encoder.set_depth(png::BitDepth::Eight);
|
|
let mut writer = encoder.write_header()?;
|
|
//let samples: &[u16] = unsafe { std::slice::from_raw_parts(p.cast(), frame.len() / 2) };
|
|
let frame = data
|
|
.iter()
|
|
.copied()
|
|
.map(|x| {
|
|
let x: f64 = x.into();
|
|
let x = x / 256.0;
|
|
((-1.665884e-08) * x.powf(4.)
|
|
+ (1.347094e-05) * x.powf(3.)
|
|
+ (-4.396264e-03) * x.powf(2.)
|
|
+ (9.506939e-01) * x
|
|
+ (-6.353247e+01)) as u8
|
|
})
|
|
.collect::<Vec<u8>>();
|
|
writer.write_image_data(&frame)?;
|
|
Ok(())
|
|
}
|
|
|
|
#[pyclass]
|
|
struct PyFrameIterator {
|
|
iter: Decoder,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl PyFrameIterator {
|
|
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
|
slf
|
|
}
|
|
fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<PyObject> {
|
|
match slf.iter.next() {
|
|
Some(f) => {
|
|
let pixels = f.pixels();
|
|
let p = Python::with_gil(|py| pixels.into_py(py));
|
|
Some(p)
|
|
}
|
|
None => None.into(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[pyfunction]
|
|
fn decode(filename: &str) -> PyResult<PyFrameIterator> {
|
|
let decoder = Decoder::new(filename).map_err(|_| PyValueError::new_err("failed to read"))?;
|
|
let iter = PyFrameIterator { iter: decoder };
|
|
Ok(iter.into())
|
|
}
|
|
|
|
/// writes to frames/<basename of filename>
|
|
pub fn decode_to_files(filename: &str) -> anyhow::Result<()> {
|
|
let frameiter = Decoder::new(filename)?;
|
|
let basename = std::path::Path::new(filename)
|
|
.file_stem()
|
|
.ok_or(anyhow::anyhow!("cannot get basename"))?
|
|
.to_str()
|
|
.ok_or(anyhow::anyhow!("cannot convert to utf-8 from os name"))?;
|
|
let target_dir = format!("frames/{}", basename);
|
|
let target_dir = std::path::Path::new(&target_dir);
|
|
if !target_dir.exists() {
|
|
std::fs::create_dir(target_dir)?;
|
|
}
|
|
for (i, frame) in frameiter.enumerate() {
|
|
let name = format!("frames/{}/{:05}.png", basename, i);
|
|
if let Err(_e) = write_raw_frame(&name, &frame.raw) {
|
|
println!("skipping bad frame {}", i);
|
|
continue;
|
|
}
|
|
let name = format!("{}/temp_{:05}.png", target_dir.display(), i);
|
|
let pixels = frame.pixels();
|
|
write_calibrated_frame(&name, &pixels)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[pymodule]
|
|
fn thermaldecoder(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
|
|
m.add_function(wrap_pyfunction!(decode, m)?)?;
|
|
Ok(())
|
|
}
|