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 Ok(PipelineData::ByteStream(stream, metadata))
80 } else {
81 operate(
82 map_value,
83 Arguments { range, cell_paths },
84 input,
85 call.head,
86 engine_state.signals(),
87 )
88 }
89 }
90
91 fn examples(&self) -> Vec<Example> {
92 vec![
93 Example {
94 description: "Extract bytes starting from a specific index",
95 example: "{ data: 0x[33 44 55 10 01 13 10] } | bytes at 3.. data",
96 result: Some(Value::test_record(record! {
97 "data" => Value::test_binary(vec![0x10, 0x01, 0x13, 0x10]),
98 })),
99 },
100 Example {
101 description: "Slice out `0x[10 01 13]` from `0x[33 44 55 10 01 13]`",
102 example: "0x[33 44 55 10 01 13] | bytes at 3..5",
103 result: Some(Value::test_binary(vec![0x10, 0x01, 0x13])),
104 },
105 Example {
106 description: "Extract bytes from the start up to a specific index",
107 example: "0x[33 44 55 10 01 13 10] | bytes at ..4",
108 result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10, 0x01])),
109 },
110 Example {
111 description: "Extract byte `0x[10]` using an exclusive end index",
112 example: "0x[33 44 55 10 01 13 10] | bytes at 3..<4",
113 result: Some(Value::test_binary(vec![0x10])),
114 },
115 Example {
116 description: "Extract bytes up to a negative index (inclusive)",
117 example: "0x[33 44 55 10 01 13 10] | bytes at ..-4",
118 result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10])),
119 },
120 Example {
121 description: "Slice bytes across multiple table columns",
122 example: r#"[[ColA ColB ColC]; [0x[11 12 13] 0x[14 15 16] 0x[17 18 19]]] | bytes at 1.. ColB ColC"#,
123 result: Some(Value::test_list(vec![Value::test_record(record! {
124 "ColA" => Value::test_binary(vec![0x11, 0x12, 0x13]),
125 "ColB" => Value::test_binary(vec![0x15, 0x16]),
126 "ColC" => Value::test_binary(vec![0x18, 0x19]),
127 })])),
128 },
129 Example {
130 description: "Extract the last three bytes using a negative start index",
131 example: "0x[33 44 55 10 01 13 10] | bytes at (-3)..",
132 result: Some(Value::test_binary(vec![0x01, 0x13, 0x10])),
133 },
134 ]
135 }
136}
137
138fn map_value(input: &Value, args: &Arguments, head: Span) -> Value {
139 let range = &args.range;
140 match input {
141 Value::Binary { val, .. } => {
142 let len = val.len() as u64;
143 let start: u64 = range.absolute_start(len);
144 let _start: usize = match start.try_into() {
145 Ok(start) => start,
146 Err(_) => {
147 let span = input.span();
148 return Value::error(
149 ShellError::UnsupportedInput {
150 msg: format!(
151 "Absolute start position {start} was too large for your system arch."
152 ),
153 input: args.range.to_string(),
154 msg_span: span,
155 input_span: span,
156 },
157 head,
158 );
159 }
160 };
161
162 let (start, end) = range.absolute_bounds(val.len());
163 let bytes: Vec<u8> = match end {
164 Bound::Unbounded => val[start..].into(),
165 Bound::Included(end) => val[start..=end].into(),
166 Bound::Excluded(end) => val[start..end].into(),
167 };
168
169 Value::binary(bytes, head)
170 }
171 Value::Error { .. } => input.clone(),
172 other => Value::error(
173 ShellError::UnsupportedInput {
174 msg: "Only binary values are supported".into(),
175 input: format!("input type: {:?}", other.get_type()),
176 msg_span: head,
177 input_span: other.span(),
178 },
179 head,
180 ),
181 }
182}
183
184#[cfg(test)]
185mod test {
186 use super::*;
187
188 #[test]
189 fn test_examples() {
190 use crate::test_examples;
191 test_examples(BytesAt {})
192 }
193}