cfn_guard/utils/
writer.rs

1use crate::rules::errors::InternalError::{
2    FromUtf8Error, IncompatibleWriterError, UnsupportedBufferError, UnsupportedOperationError,
3};
4use crate::Error;
5use std::fs::File;
6use std::io::{Read, Stderr, Stdout, Write};
7
8#[derive(Debug)]
9pub struct Writer {
10    buffer: WriteBuffer,
11    err: WriteBuffer,
12}
13
14impl Default for Writer {
15    fn default() -> Self {
16        Self {
17            buffer: WriteBuffer::Stdout(std::io::stdout()),
18            err: WriteBuffer::Stderr(std::io::stderr()),
19        }
20    }
21}
22
23impl Writer {
24    pub fn new(buffer: WriteBuffer) -> crate::rules::Result<Self> {
25        if buffer.is_err() {
26            return Err(Error::from(IncompatibleWriterError(
27                "Unable to use stderr as a regular buffer.".to_string(),
28            )));
29        }
30
31        Ok(Self {
32            buffer,
33            err: WriteBuffer::Stderr(std::io::stderr()),
34        })
35    }
36
37    pub fn new_with_err(buffer: WriteBuffer, err: WriteBuffer) -> crate::rules::Result<Self> {
38        if buffer.is_err() {
39            return Err(Error::from(IncompatibleWriterError(
40                "Unable to use stderr as a regular buffer.".to_string(),
41            )));
42        }
43
44        Ok(Self { buffer, err })
45    }
46
47    pub fn write_err(&mut self, s: String) -> std::io::Result<()> {
48        writeln!(self.err, "{s}")
49    }
50
51    pub fn err_to_stripped(self) -> crate::rules::Result<String> {
52        match self.err {
53            WriteBuffer::Vec(vec) => String::from_utf8(strip_ansi_escapes::strip(vec)?)
54                .map_err(|e| Error::from(FromUtf8Error(e))),
55            WriteBuffer::File(mut file) => {
56                let mut data = String::new();
57                file.read_to_string(&mut data)
58                    .expect("Unable to read from file");
59
60                String::from_utf8(strip_ansi_escapes::strip(data)?)
61                    .map_err(|e| Error::from(FromUtf8Error(e)))
62            }
63            WriteBuffer::Stdout(..) => Err(Error::from(UnsupportedOperationError(
64                "Unable to call err_to_stripped() on a stdout buffer.".to_string(),
65            ))),
66            WriteBuffer::Stderr(..) => Err(Error::from(UnsupportedOperationError(
67                "Unable to call err_to_stripped() on a stderr buffer.".to_string(),
68            ))),
69        }
70    }
71
72    pub fn into_string(self) -> crate::rules::Result<String> {
73        self.buffer.into_string()
74    }
75
76    pub fn stripped(self) -> crate::rules::Result<String> {
77        match self.buffer {
78            WriteBuffer::Vec(vec) => {
79                let stripped = strip_ansi_escapes::strip(vec)?;
80
81                String::from_utf8(stripped).map_err(|e| Error::from(FromUtf8Error(e)))
82            }
83            WriteBuffer::File(mut file) => {
84                let mut data = String::new();
85                file.read_to_string(&mut data)
86                    .expect("Unable to read from file");
87
88                let stripped = strip_ansi_escapes::strip(data.into_bytes())?;
89
90                String::from_utf8(stripped).map_err(|e| Error::from(FromUtf8Error(e)))
91            }
92            WriteBuffer::Stdout(..) => Err(Error::from(UnsupportedBufferError(
93                "Unable to strip ANSI escapes from stdout buffer.".to_string(),
94            ))),
95            WriteBuffer::Stderr(..) => Err(Error::from(UnsupportedBufferError(
96                "Unable to strip ANSI escapes from stderr buffer.".to_string(),
97            ))),
98        }
99    }
100}
101
102impl Write for Writer {
103    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
104        self.buffer.write(String::from_utf8_lossy(buf).as_bytes())
105    }
106
107    fn flush(&mut self) -> std::io::Result<()> {
108        self.buffer.flush()
109    }
110}
111
112#[derive(Debug)]
113pub enum WriteBuffer {
114    Stdout(Stdout),
115    Vec(Vec<u8>),
116    File(File),
117    Stderr(Stderr),
118}
119
120impl WriteBuffer {
121    fn is_err(&self) -> bool {
122        matches!(self, WriteBuffer::Stderr(_))
123    }
124    fn into_string(self) -> crate::rules::Result<String> {
125        match self {
126            WriteBuffer::Stdout(..) => Err(Error::from(UnsupportedOperationError(
127                "Unable to call into_string() on a stdout buffer.".to_string(),
128            ))),
129            WriteBuffer::Stderr(..) => Err(Error::from(UnsupportedOperationError(
130                "Unable to call into_string() on a stderr buffer.".to_string(),
131            ))),
132            WriteBuffer::Vec(vec) => {
133                String::from_utf8(vec).map_err(|e| Error::from(FromUtf8Error(e)))
134            }
135            WriteBuffer::File(mut file) => {
136                let mut data = String::new();
137                file.read_to_string(&mut data)
138                    .expect("Unable to read from file");
139
140                Ok(data)
141            }
142        }
143    }
144}
145
146impl Write for WriteBuffer {
147    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
148        match self {
149            WriteBuffer::Stdout(stdout) => stdout.write(buf),
150            WriteBuffer::Stderr(stderr) => stderr.write(buf),
151            WriteBuffer::Vec(vec) => vec.write(buf),
152            WriteBuffer::File(file) => file.write(buf),
153        }
154    }
155
156    fn flush(&mut self) -> std::io::Result<()> {
157        match self {
158            WriteBuffer::Stdout(stdout) => stdout.flush(),
159            WriteBuffer::Stderr(stderr) => stderr.flush(),
160            WriteBuffer::Vec(vec) => vec.flush(),
161            WriteBuffer::File(file) => file.flush(),
162        }
163    }
164}
cfn_guard/utils/writer.rs

cfn_guard/utils/
writer.rs
