log2db/src/main.rs
2019-10-10 14:39:41 +02:00

150 lines
4.9 KiB
Rust

use std::str::FromStr;
use tokio_postgres::{Config, NoTls, Client, Statement};
use tokio::net::{TcpListener, TcpStream};
use tokio::prelude::*;
use std::net::SocketAddr;
use std::sync::Arc;
use structopt::StructOpt;
use std::collections::HashSet;
type Error = Box<dyn std::error::Error>;
/// This application recieves logfiles in ubiquitis weird logfile format from the network,
/// splits each line into its fields and writes them to a sql database for further analysis.
///
/// The main required argument is an postgres connection string, passed in the environment
/// variable LOG2DB_PSQL
#[derive(StructOpt)]
struct Options {
/// ip and (TCP) port to bind to
#[structopt(short, long, default_value="[::]:514")]
addr : String,
/// space separated list of daemon/service names to not write to the db
#[structopt(default_value="dnsmasq-dhcp")]
blacklist: Vec<String>
}
#[tokio::main]
async fn main() -> Result<(), Error> {
let options = Options::from_args();
let postgresurl = std::env::var("LOG2DB_PSQL")
.map_err(|_| "please supply the LOG2DB_PSQL environment variable, containting a postgress connection string")?;
let mut cfg = Config::from_str(&postgresurl)?;
cfg.application_name("log2db");
let (client, connection) = cfg.connect(NoTls).await?;
let connection = connection.map(|r| {
if let Err(e) = r {
eprintln!("could not connect to database: {}", e);
}
});
tokio::spawn(connection);
client.execute("create table if not exists log(
prio smallint,
rcv_ip inet,
rcv_date timestamptz,
date timestamptz,
daemon varchar,
message varchar
)
", &[]).await?;
let client = Arc::new(client);
let insert_statement = client
.prepare("insert into log(prio, rcv_ip, rcv_date, date, daemon, message) values ($1, $2, $3, $4, $5, $6)").await?;
let mut listener = TcpListener::bind(&options.addr).await
.map_err(|e| format!("could not bind to {} with error {:?}", &options.addr, e))?;
let blacklist: HashSet<_> = options.blacklist.into_iter().collect();
let blacklist = Arc::new(blacklist);
loop {
match listener.accept().await {
Ok((socket, peer)) => {
tokio::spawn(
handle_peer_and_error(socket, peer, client.clone(), insert_statement.clone(), blacklist.clone())
);
},
Err(e) => eprintln!("{:?}", e),
}
}
}
use chrono::{Local, DateTime, FixedOffset};
async fn handle_peer_and_error(stream: TcpStream, peer: SocketAddr, db: Arc<Client>, insert_statement: Statement, blacklist: Arc<HashSet<String>>) {
if let Err(e) = handle_peer(stream, peer, db, insert_statement, blacklist).await {
eprintln!("{}", e);
}
}
async fn handle_peer(stream: TcpStream, peer: SocketAddr, db: Arc<Client>, insert_statement: Statement, blacklist: Arc<HashSet<String>>)
-> Result<(), Error> {
use tokio::codec::{FramedRead, LinesCodec};
let ip = peer.ip();
let mut lines = FramedRead::new(stream, LinesCodec::new());
loop {
match lines.next().await.transpose()? {
Some(line) => {
let (prio, now, date, service, log) = parse_line(&line)?;
if !blacklist.contains(service) {
db.execute(&insert_statement, &[&prio, &ip, &now, &date, &service, &log]).await?;
}
},
None => break,
}
}
Ok(())
}
/** parses a line, returning
* ( time the log was recieved
* , time the log was written according to logger
* , name of the service that wrote the log
* , log entry
* )
*/
fn parse_line(line: &'_ str) -> Result<(i16, DateTime<Local>, DateTime<FixedOffset>, &'_ str, &'_ str), Error> {
let mut prio_and_remainder = line.splitn(2, '>');
let prio = prio_and_remainder.next().ok_or("log did not contain priority")?;
let prio = &prio[1..];
let prio = prio.parse()
.map_err(|e| format!("could not parse priority {}: {}",prio, e))?;
let line = prio_and_remainder.next().expect("splitn should always return a second part");
let (date, line) = line.split_at(16);
// we need to prepend the current year and timezone, as that is not stated in the logfile
let now = chrono::Local::now();
let mut base = format!("{}", now.format("%Y %z "));
base.push_str(date);
let date = DateTime::parse_from_str(&base, "%Y %z %b %e %H:%M:%S ").
map_err(|e| format!("could not parse {}{} {}", date, line, e))?;
let mut parts = line.splitn(2, ':');
let service_and_pid = parts.next().ok_or("could not parse service")?;
let mut service_parts = service_and_pid.splitn(2, '[');
let service = service_parts.next().ok_or("could not split pid from service")?.trim();
let log = parts.next().ok_or("could not parse logfile")?.trim();
Ok((prio, now, date, service, log))
}