1use std::ops::Bound;
2
3use nu_cmd_base::input_handler::{CmdArgument, operate};
4use nu_engine::command_prelude::*;
5use nu_protocol::{IntRange, Range};
6
7#[derive(Clone)]
8pub struct BytesAt;
9
10struct Arguments {
11 range: IntRange,
12 cell_paths: Option<Vec<CellPath>>,
13}
14
15impl CmdArgument for Arguments {
16 fn take_cell_paths(&mut self) -> Option<Vec<CellPath>> {
17 self.cell_paths.take()
18 }
19}
20
21impl Command for BytesAt {
22 fn name(&self) -> &str {
23 "bytes at"
24 }
25
26 fn signature(&self) -> Signature {
27 Signature::build("bytes at")
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("range", SyntaxShape::Range, "The range to get bytes.")
39 .rest(
40 "rest",
41 SyntaxShape::CellPath,
42 "For a data structure input, get bytes from data at the given cell paths.",
43 )
44 .category(Category::Bytes)
45 }
46
47 fn description(&self) -> &str {
48 "Get bytes defined by a range."
49 }
50
51 fn search_terms(&self) -> Vec<&str> {
52 vec!["slice"]
53 }
54
55 fn run(
56 &self,
57 engine_state: &EngineState,
58 stack: &mut Stack,
59 call: &Call,
60 input: PipelineData,
61 ) -> Result<PipelineData, ShellError> {
62 let range = match call.req(engine_state, stack, 0)? {
63 Range::IntRange(range) => range,
64 _ => {
65 return Err(ShellError::UnsupportedInput {
66 msg: "Float ranges are not supported for byte streams".into(),
67 input: "value originates from here".into(),
68 msg_span: call.head,
69 input_span: call.head,
70 });
71 }
72 };
73
74 let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
75 let cell_paths = (!cell_paths.is_empty()).then_some(cell_paths);
76
77 if let PipelineData::ByteStream(stream, metadata) = input {
78 let stream = stream.slice(call.head, call.arguments_span(), range)?;
79 let metadata = metadata.map(|m| m.with_content_type(None));
81 Ok(PipelineData::byte_stream(stream, metadata))
82 } else {
83 operate(
84 map_value,
85 Arguments { range, cell_paths },
86 input,
87 call.head,
88 engine_state.signals(),
89 )
90 .map(|pipeline| {
91 let metadata = pipeline.metadata().map(|m| m.with_content_type(None));
93 pipeline.set_metadata(metadata)
94 })
95 }
96 }
97
98 fn examples(&self) -> Vec<Example> {
99 vec![
100 Example {
101 description: "Extract bytes starting from a specific index",
102 example: "{ data: 0x[33 44 55 10 01 13 10] } | bytes at 3.. data",
103 result: Some(Value::test_record(record! {
104 "data" => Value::test_binary(vec![0x10, 0x01, 0x13, 0x10]),
105 })),
106 },
107 Example {
108 description: "Slice out `0x[10 01 13]` from `0x[33 44 55 10 01 13]`",
109 example: "0x[33 44 55 10 01 13] | bytes at 3..5",
110 result: Some(Value::test_binary(vec![0x10, 0x01, 0x13])),
111 },
112 Example {
113 description: "Extract bytes from the start up to a specific index",
114 example: "0x[33 44 55 10 01 13 10] | bytes at ..4",
115 result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10, 0x01])),
116 },
117 Example {
118 description: "Extract byte `0x[10]` using an exclusive end index",
119 example: "0x[33 44 55 10 01 13 10] | bytes at 3..<4",
120 result: Some(Value::test_binary(vec![0x10])),
121 },
122 Example {
123 description: "Extract bytes up to a negative index (inclusive)",
124 example: "0x[33 44 55 10 01 13 10] | bytes at ..-4",
125 result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10])),
126 },
127 Example {
128 description: "Slice bytes across multiple table columns",
129 example: r#"[[ColA ColB ColC]; [0x[11 12 13] 0x[14 15 16] 0x[17 18 19]]] | bytes at 1.. ColB ColC"#,
130 result: Some(Value::test_list(vec![Value::test_record(record! {
131 "ColA" => Value::test_binary(vec![0x11, 0x12, 0x13]),
132 "ColB" => Value::test_binary(vec![0x15, 0x16]),
133 "ColC" => Value::test_binary(vec![0x18, 0x19]),
134 })])),
135 },
136 Example {
137 description: "Extract the last three bytes using a negative start index",
138 example: "0x[33 44 55 10 01 13 10] | bytes at (-3)..",
139 result: Some(Value::test_binary(vec![0x01, 0x13, 0x10])),
140 },
141 ]
142 }
143}
144
145fn map_value(input: &Value, args: &Arguments, head: Span) -> Value {
146 let range = &args.range;
147 match input {
148 Value::Binary { val, .. } => {
149 let len = val.len() as u64;
150 let start: u64 = range.absolute_start(len);
151 let _start: usize = match start.try_into() {
152 Ok(start) => start,
153 Err(_) => {
154 let span = input.span();
155 return Value::error(
156 ShellError::UnsupportedInput {
157 msg: format!(
158 "Absolute start position {start} was too large for your system arch."
159 ),
160 input: args.range.to_string(),
161 msg_span: span,
162 input_span: span,
163 },
164 head,
165 );
166 }
167 };
168
169 let (start, end) = range.absolute_bounds(val.len());
170 let bytes: Vec<u8> = match end {
171 Bound::Unbounded => val[start..].into(),
172 Bound::Included(end) => val[start..=end].into(),
173 Bound::Excluded(end) => val[start..end].into(),
174 };
175
176 Value::binary(bytes, head)
177 }
178 Value::Error { .. } => input.clone(),
179 other => Value::error(
180 ShellError::UnsupportedInput {
181 msg: "Only binary values are supported".into(),
182 input: format!("input type: {:?}", other.get_type()),
183 msg_span: head,
184 input_span: other.span(),
185 },
186 head,
187 ),
188 }
189}
190
191#[cfg(test)]
192mod test {
193 use super::*;
194
195 #[test]
196 fn test_examples() {
197 use crate::test_examples;
198 test_examples(BytesAt {})
199 }
200}