186 lines
8.3 KiB
Rust
186 lines
8.3 KiB
Rust
extern crate tokio_postgres;
|
|
extern crate serde_json;
|
|
extern crate tokio;
|
|
use serde::ser::{Serialize, Serializer, SerializeStruct};
|
|
mod parser;
|
|
use std::fs::File;
|
|
use serde_json::json;
|
|
|
|
use tokio_postgres::types::ToSql;
|
|
//use futures::{TryStreamExt};
|
|
use tokio_postgres::{NoTls, Error};
|
|
|
|
|
|
impl Serialize for parser::QryData {
|
|
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
|
where
|
|
S: Serializer,
|
|
{
|
|
// 34 is the number of fields in the struct.
|
|
let mut state = serializer.serialize_struct("parser::QryData", 34)?;
|
|
state.serialize_field("time", &self.time)?;
|
|
state.serialize_field("ether_header.ether_dhost", &self.ether_header.ether_dhost)?;
|
|
state.serialize_field("ether_header.ether_shost", &self.ether_header.ether_shost)?;
|
|
state.serialize_field("ether_header.ether_type", &self.ether_header.ether_type)?;
|
|
state.serialize_field("ipv4_header.ip_version", &self.ipv4_header.ip_version)?;
|
|
state.serialize_field("ipv4_header.ip_ihl", &self.ipv4_header.ip_ihl)?;
|
|
state.serialize_field("ipv4_header.ip_dscp", &self.ipv4_header.ip_dscp)?;
|
|
state.serialize_field("ipv4_header.ip_ecn", &self.ipv4_header.ip_ecn)?;
|
|
state.serialize_field("ipv4_header.ip_total_length", &self.ipv4_header.ip_total_length)?;
|
|
state.serialize_field("ipv4_header.ip_identification", &self.ipv4_header.ip_identification)?;
|
|
state.serialize_field("ipv4_header.ip_df", &self.ipv4_header.ip_df)?;
|
|
state.serialize_field("ipv4_header.ip_mf", &self.ipv4_header.ip_mf)?;
|
|
state.serialize_field("ipv4_header.ip_fragment_offset", &self.ipv4_header.ip_fragment_offset)?;
|
|
state.serialize_field("ipv4_header.ip_source_address", &self.ipv4_header.ip_source_address)?;
|
|
state.serialize_field("ipv4_header.ip_destination_address", &self.ipv4_header.ip_destination_address)?;
|
|
state.serialize_field("tcp_header.source_port", &self.tcp_header.source_port)?;
|
|
state.serialize_field("tcp_header.destination_port", &self.tcp_header.destination_port)?;
|
|
state.serialize_field("tcp_header.seq_num", &self.tcp_header.seq_num)?;
|
|
state.serialize_field("tcp_header.ack_num", &self.tcp_header.ack_num)?;
|
|
state.serialize_field("tcp_header.data_offset", &self.tcp_header.data_offset)?;
|
|
state.serialize_field("tcp_header.reserved", &self.tcp_header.reserved)?;
|
|
state.serialize_field("tcp_header.ns", &self.tcp_header.ns)?;
|
|
state.serialize_field("tcp_header.cwr", &self.tcp_header.cwr)?;
|
|
state.serialize_field("tcp_header.ece", &self.tcp_header.ece)?;
|
|
state.serialize_field("tcp_header.urg", &self.tcp_header.urg)?;
|
|
state.serialize_field("tcp_header.ack", &self.tcp_header.ack)?;
|
|
state.serialize_field("tcp_header.psh", &self.tcp_header.psh)?;
|
|
state.serialize_field("tcp_header.rst", &self.tcp_header.rst)?;
|
|
state.serialize_field("tcp_header.syn", &self.tcp_header.syn)?;
|
|
state.serialize_field("tcp_header.fin", &self.tcp_header.fin)?;
|
|
state.serialize_field("tcp_header.window_size", &self.tcp_header.window_size)?;
|
|
state.serialize_field("tcp_header.checksum", &self.tcp_header.checksum)?;
|
|
state.serialize_field("tcp_header.urgent_pointer", &self.tcp_header.urgent_pointer)?;
|
|
state.serialize_field("data", &self.data)?;
|
|
state.end()
|
|
}
|
|
}
|
|
|
|
fn serialize_packets ( v: Vec<parser::QryData> ) -> Vec<serde_json::Value> {
|
|
let mut packets_serialized: Vec<_> = Vec::new();
|
|
|
|
for packet in v.iter() {
|
|
packets_serialized.push(json!(&packet));
|
|
}
|
|
packets_serialized
|
|
}
|
|
|
|
fn query_string ( insert_max: &usize ) -> String {
|
|
let mut insert_template: String = "INSERT INTO json_dump (packet) Values ".to_owned();
|
|
|
|
for insert in 0..insert_max-1 { insert_template.push_str( &(format!("(${}), ", insert+1)) );}
|
|
insert_template.push_str( &(format!("(${})", insert_max)) );
|
|
|
|
insert_template
|
|
}
|
|
|
|
#[tokio::main(core_threads = 4)] // By default, tokio_postgres uses the tokio crate as its runtime.
|
|
async fn main() -> Result<(), Error> {
|
|
|
|
/* Init values from file */
|
|
let file = File::open("parser.json").expect("file should open read only");
|
|
|
|
let json: serde_json::Value = serde_json::from_reader(file).unwrap();
|
|
let filter = json.get("filter").unwrap().as_str().unwrap();
|
|
let insert_max = json.get("insert_max").unwrap().as_u64().unwrap() as usize;
|
|
let pcap_file = json.get("pcap_file").unwrap().as_str().unwrap();
|
|
let host = ["host=", json.get("database_host").unwrap().as_str().unwrap()].join("");
|
|
let user = ["user=", json.get("database_user").unwrap().as_str().unwrap()].join("");
|
|
let password = ["password=", json.get("database_password").unwrap().as_str().unwrap()].join("");
|
|
let connection = [host, user, password].join(" ");
|
|
let device = json.get("parse_device").unwrap().as_str().unwrap();
|
|
let is_device = json.get("from_device").unwrap().as_bool().unwrap();
|
|
|
|
|
|
/* db connection */
|
|
let (client, connection) =
|
|
tokio_postgres::connect(&connection, NoTls).await?;
|
|
|
|
tokio::spawn(async move {
|
|
if let Err(e) = connection.await {
|
|
eprintln!("connection error: {}", e);
|
|
}
|
|
});
|
|
|
|
client.execute("DROP TABLE IF EXISTS json_dump", &[]).await?;
|
|
client.execute("CREATE TABLE json_dump ( ID serial NOT NULL PRIMARY KEY, packet json NOT NULL)", &[]).await?;
|
|
|
|
/* device or file input */
|
|
if false == is_device {
|
|
|
|
let v: Vec<parser::QryData> = parser::parse(&pcap_file, &filter );
|
|
let packets_serialized = serialize_packets( v );
|
|
|
|
/* Query */
|
|
//let insert_max = 60;
|
|
let chunk_count = packets_serialized.len()/insert_max;
|
|
let remainder: usize = packets_serialized.len() % insert_max;
|
|
let chunker = &packets_serialized.len() < &insert_max;
|
|
match chunker {
|
|
true => {
|
|
let insert_str = query_string( &packets_serialized.len() );
|
|
let statement_false = client.prepare( &insert_str ).await?;
|
|
client.query_raw(&statement_false, packets_serialized.iter().map(|p| p as &dyn ToSql)).await?;
|
|
},
|
|
|
|
false => {
|
|
|
|
let insert_str = query_string( &insert_max );
|
|
let statement = client.prepare( &insert_str ).await?;
|
|
|
|
|
|
for _i in 0..chunk_count {
|
|
let (_input, _)= packets_serialized.split_at(insert_max);
|
|
client.query_raw(&statement, _input.to_vec().iter().map(|p| p as &dyn ToSql)).await?;
|
|
}
|
|
|
|
println!("Packets, total:{:?}",packets_serialized.len());
|
|
println!("Chunks, total:{}", chunk_count);
|
|
println!("Chunks, remainder{}", remainder);
|
|
|
|
if remainder > 0 {
|
|
let rem_str = query_string( &remainder );
|
|
let statement_remainder = client.prepare(&rem_str).await?;
|
|
let (_garbage, _input) =packets_serialized.split_at(packets_serialized.len()-remainder);
|
|
client.query_raw(&statement_remainder, _input.to_vec().iter().map(|p| p as &dyn ToSql),).await?;
|
|
}
|
|
|
|
}
|
|
}
|
|
|
|
} else {
|
|
let insert_str = query_string(&insert_max);
|
|
let statement = client.prepare(&insert_str).await?;
|
|
loop {
|
|
let v: Vec<parser::QryData> = parser::parse_device(&device, &filter, &insert_max);
|
|
let packets_serialized = serialize_packets( v );
|
|
client.query_raw(&statement, packets_serialized.iter().map(|p| p as &dyn ToSql),).await?;
|
|
}
|
|
}
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
#[test]
|
|
fn test_insert_json () {
|
|
use serde_json::json;
|
|
let mut client = Client::connect("host=localhost user=postgres password=password", NoTls).unwrap();
|
|
let john = json!({
|
|
"name": "John Doe",
|
|
"age": 43,
|
|
"phones": [
|
|
"+44 1234567",
|
|
"+44 2345678"
|
|
],
|
|
"empty": []
|
|
});
|
|
|
|
client.execute("DROP TABLE IF EXISTS json_dump", &[]).unwrap();
|
|
client.execute("CREATE TABLE json_dump ( ID serial NOT NULL PRIMARY KEY, data json NOT NULL)", &[]);
|
|
client.query("INSERT INTO json_dump ( data ) VALUES ($1)", &[&john]);
|
|
}
|
|
|
|
|