nu_command/bytes/
add.rs

1use nu_cmd_base::input_handler::{CmdArgument, operate};
2use nu_engine::command_prelude::*;
3
4struct Arguments {
5    added_data: Vec<u8>,
6    index: Option<usize>,
7    end: bool,
8    cell_paths: Option<Vec<CellPath>>,
9}
10
11impl CmdArgument for Arguments {
12    fn take_cell_paths(&mut self) -> Option<Vec<CellPath>> {
13        self.cell_paths.take()
14    }
15}
16
17#[derive(Clone)]
18
19pub struct BytesAdd;
20
21impl Command for BytesAdd {
22    fn name(&self) -> &str {
23        "bytes add"
24    }
25
26    fn signature(&self) -> Signature {
27        Signature::build("bytes add")
28            .input_output_types(vec![
29                (Type::Binary, Type::Binary),
30                (
31                    Type::List(Box::new(Type::Binary)),
32                    Type::List(Box::new(Type::Binary)),
33                ),
34                (Type::table(), Type::table()),
35                (Type::record(), Type::record()),
36            ])
37            .allow_variants_without_examples(true)
38            .required("data", SyntaxShape::Binary, "The binary to add.")
39            .named(
40                "index",
41                SyntaxShape::Int,
42                "index to insert binary data",
43                Some('i'),
44            )
45            .switch("end", "add to the end of binary", Some('e'))
46            .rest(
47                "rest",
48                SyntaxShape::CellPath,
49                "For a data structure input, add bytes to the data at the given cell paths.",
50            )
51            .category(Category::Bytes)
52    }
53
54    fn description(&self) -> &str {
55        "Add specified bytes to the input."
56    }
57
58    fn search_terms(&self) -> Vec<&str> {
59        vec!["append", "truncate", "padding"]
60    }
61
62    fn run(
63        &self,
64        engine_state: &EngineState,
65        stack: &mut Stack,
66        call: &Call,
67        input: PipelineData,
68    ) -> Result<PipelineData, ShellError> {
69        let added_data: Vec<u8> = call.req(engine_state, stack, 0)?;
70        let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
71        let cell_paths = (!cell_paths.is_empty()).then_some(cell_paths);
72        let index: Option<usize> = call.get_flag(engine_state, stack, "index")?;
73        let end = call.has_flag(engine_state, stack, "end")?;
74
75        let arg = Arguments {
76            added_data,
77            index,
78            end,
79            cell_paths,
80        };
81        operate(add, arg, input, call.head, engine_state.signals())
82    }
83
84    fn examples(&self) -> Vec<Example> {
85        vec![
86            Example {
87                description: "Add bytes `0x[AA]` to `0x[1F FF AA AA]`",
88                example: "0x[1F FF AA AA] | bytes add 0x[AA]",
89                result: Some(Value::binary(
90                    vec![0xAA, 0x1F, 0xFF, 0xAA, 0xAA],
91                    Span::test_data(),
92                )),
93            },
94            Example {
95                description: "Add bytes `0x[AA BB]` to `0x[1F FF AA AA]` at index 1",
96                example: "0x[1F FF AA AA] | bytes add 0x[AA BB] --index 1",
97                result: Some(Value::binary(
98                    vec![0x1F, 0xAA, 0xBB, 0xFF, 0xAA, 0xAA],
99                    Span::test_data(),
100                )),
101            },
102            Example {
103                description: "Add bytes `0x[11]` to `0x[FF AA AA]` at the end",
104                example: "0x[FF AA AA] | bytes add 0x[11] --end",
105                result: Some(Value::binary(
106                    vec![0xFF, 0xAA, 0xAA, 0x11],
107                    Span::test_data(),
108                )),
109            },
110            Example {
111                description: "Add bytes `0x[11 22 33]` to `0x[FF AA AA]` at the end, at index 1(the index is start from end)",
112                example: "0x[FF AA BB] | bytes add 0x[11 22 33] --end --index 1",
113                result: Some(Value::binary(
114                    vec![0xFF, 0xAA, 0x11, 0x22, 0x33, 0xBB],
115                    Span::test_data(),
116                )),
117            },
118        ]
119    }
120}
121
122fn add(val: &Value, args: &Arguments, span: Span) -> Value {
123    let val_span = val.span();
124    match val {
125        Value::Binary { val, .. } => add_impl(val, args, val_span),
126        // Propagate errors by explicitly matching them before the final case.
127        Value::Error { .. } => val.clone(),
128        other => Value::error(
129            ShellError::OnlySupportsThisInputType {
130                exp_input_type: "binary".into(),
131                wrong_type: other.get_type().to_string(),
132                dst_span: span,
133                src_span: other.span(),
134            },
135            span,
136        ),
137    }
138}
139
140fn add_impl(input: &[u8], args: &Arguments, span: Span) -> Value {
141    match args.index {
142        None => {
143            if args.end {
144                let mut added_data = args.added_data.clone();
145                let mut result = input.to_vec();
146                result.append(&mut added_data);
147                Value::binary(result, span)
148            } else {
149                let mut result = args.added_data.clone();
150                let mut input = input.to_vec();
151                result.append(&mut input);
152                Value::binary(result, span)
153            }
154        }
155        Some(mut indx) => {
156            let inserted_index = if args.end {
157                input.len().saturating_sub(indx)
158            } else {
159                if indx > input.len() {
160                    indx = input.len()
161                }
162                indx
163            };
164            let mut result = vec![];
165            let mut prev_data = input[..inserted_index].to_vec();
166            result.append(&mut prev_data);
167            let mut added_data = args.added_data.clone();
168            result.append(&mut added_data);
169            let mut after_data = input[inserted_index..].to_vec();
170            result.append(&mut after_data);
171            Value::binary(result, span)
172        }
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn test_examples() {
182        use crate::test_examples;
183
184        test_examples(BytesAdd {})
185    }
186}