use std::str;
use std;
use nom7::IResult;
use nom7::combinator::map_res;
use nom7::bytes::streaming::{tag, take_while};
use nom7::number::streaming::be_u8;
use crate::applayer::{AppLayerTxData,AppLayerStateData};
const READREQUEST: u8 = 1;
const WRITEREQUEST: u8 = 2;
const DATA: u8 = 3;
const ACK: u8 = 4;
const ERROR: u8 = 5;
#[derive(Debug, PartialEq, Eq)]
pub struct TFTPTransaction {
pub opcode : u8,
pub filename : String,
pub mode : String,
id: u64,
tx_data: AppLayerTxData,
}
pub struct TFTPState {
state_data: AppLayerStateData,
pub transactions : Vec<TFTPTransaction>,
tx_id: u64,
}
impl TFTPState {
fn get_tx_by_id(&mut self, tx_id: u64) -> Option<&TFTPTransaction> {
self.transactions.iter().find(|&tx| tx.id == tx_id + 1)
}
fn free_tx(&mut self, tx_id: u64) {
let tx = self.transactions.iter().position(|tx| tx.id == tx_id + 1);
debug_assert!(tx.is_some());
if let Some(idx) = tx {
let _ = self.transactions.remove(idx);
}
}
}
impl TFTPTransaction {
pub fn new(opcode : u8, filename : String, mode : String) -> TFTPTransaction {
TFTPTransaction {
opcode,
filename,
mode : mode.to_lowercase(),
id : 0,
tx_data: AppLayerTxData::new(),
}
}
pub fn is_mode_ok(&self) -> bool {
match self.mode.as_str() {
"netascii" | "mail" | "octet" => true,
_ => false
}
}
pub fn is_opcode_ok(&self) -> bool {
match self.opcode {
READREQUEST | WRITEREQUEST | ACK | DATA | ERROR => true,
_ => false
}
}
}
#[no_mangle]
pub extern "C" fn SCTftpStateAlloc() -> *mut std::os::raw::c_void {
let state = TFTPState { state_data: AppLayerStateData::new(), transactions : Vec::new(), tx_id: 0, };
let boxed = Box::new(state);
return Box::into_raw(boxed) as *mut _;
}
#[no_mangle]
pub extern "C" fn SCTftpStateFree(state: *mut std::os::raw::c_void) {
std::mem::drop(unsafe { Box::from_raw(state as *mut TFTPState) });
}
#[no_mangle]
pub extern "C" fn SCTftpStateTxFree(state: &mut TFTPState,
tx_id: u64) {
state.free_tx(tx_id);
}
#[no_mangle]
pub extern "C" fn SCTftpGetTx(state: &mut TFTPState,
tx_id: u64) -> *mut std::os::raw::c_void {
match state.get_tx_by_id(tx_id) {
Some(tx) => tx as *const _ as *mut _,
None => std::ptr::null_mut(),
}
}
#[no_mangle]
pub extern "C" fn SCTftpGetTxCnt(state: &mut TFTPState) -> u64 {
return state.tx_id;
}
fn getstr(i: &[u8]) -> IResult<&[u8], &str> {
map_res(
take_while(|c| c != 0),
str::from_utf8
)(i)
}
fn tftp_request(slice: &[u8]) -> IResult<&[u8], TFTPTransaction> {
let (i, _) = tag([0])(slice)?;
let (i, opcode) = be_u8(i)?;
let (i, filename) = getstr(i)?;
let (i, _) = tag([0])(i)?;
let (i, mode) = getstr(i)?;
Ok((i,
TFTPTransaction::new(opcode, String::from(filename), String::from(mode))
)
)
}
fn parse_tftp_request(input: &[u8]) -> Option<TFTPTransaction> {
match tftp_request(input) {
Ok((_, tx)) => {
if !tx.is_mode_ok() {
return None;
}
if !tx.is_opcode_ok() {
return None;
}
return Some(tx);
}
Err(_) => {
return None;
}
}
}
#[no_mangle]
pub unsafe extern "C" fn SCTftpParseRequest(state: &mut TFTPState,
input: *const u8,
len: u32) -> i64 {
let buf = std::slice::from_raw_parts(input, len as usize);
match parse_tftp_request(buf) {
Some(mut tx) => {
state.tx_id += 1;
tx.id = state.tx_id;
state.transactions.push(tx);
0
},
None => {
-1
}
}
}
#[no_mangle]
pub unsafe extern "C" fn SCTftpGetTxData(
tx: *mut std::os::raw::c_void)
-> *mut AppLayerTxData
{
let tx = cast_pointer!(tx, TFTPTransaction);
return &mut tx.tx_data;
}
#[no_mangle]
pub unsafe extern "C" fn SCTftpGetStateData(
state: *mut std::os::raw::c_void)
-> *mut AppLayerStateData
{
let state = cast_pointer!(state, TFTPState);
return &mut state.state_data;
}
#[cfg(test)]
mod test {
use super::*;
static READ_REQUEST: [u8; 20] = [
0x00, 0x01, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
];
static READ_REQUEST_INVALID_1: [u8; 20] = [
0x00, 0x01, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x6e, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
];
static READ_REQUEST_INVALID_2: [u8; 3] = [
0xff, 0xff, 0xff,
];
static WRITE_REQUEST: [u8; 20] = [
0x00, 0x02, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
];
static INVALID_OPCODE: [u8; 20] = [
0x00, 0x06, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x6e, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
];
static INVALID_MODE: [u8; 20] = [
0x00, 0x01, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x63, 0x63, 0x63, 0x63, 0x63, 0x00,
];
#[test]
pub fn test_parse_tftp_read_request_1() {
let tx = TFTPTransaction {
opcode: READREQUEST,
filename: String::from("rfc1350.txt"),
mode: String::from("octet"),
id: 0,
tx_data: AppLayerTxData::new(),
};
let txp = parse_tftp_request(&READ_REQUEST[..]).unwrap();
assert_eq!(tx, txp);
}
#[test]
pub fn test_parse_tftp_write_request_1() {
let tx = TFTPTransaction {
opcode: WRITEREQUEST,
filename: String::from("rfc1350.txt"),
mode: String::from("octet"),
id: 0,
tx_data: AppLayerTxData::new(),
};
let txp = parse_tftp_request(&WRITE_REQUEST[..]).unwrap();
assert_eq!(tx, txp);
}
#[test]
pub fn test_parse_tftp_read_request_2() {
assert_eq!(None, parse_tftp_request(&READ_REQUEST_INVALID_1[..]));
}
#[test]
pub fn test_parse_tftp_read_request_3() {
assert_eq!(None, parse_tftp_request(&READ_REQUEST_INVALID_2[..]));
}
#[test]
pub fn test_parse_tftp_invalid_opcode_1() {
assert_eq!(None, parse_tftp_request(&INVALID_OPCODE[..]));
}
#[test]
pub fn test_parse_tftp_invalid_mode() {
assert_eq!(None, parse_tftp_request(&INVALID_MODE[..]));
}
}