use crate::core::types::{ProtocolMatch, RefractiumProtocol, Transport};
use memchr::memmem;
use std::sync::Arc;
#[cfg(feature = "proto-dns")]
pub mod dns;
#[cfg(feature = "proto-ftp")]
pub mod ftp;
#[cfg(feature = "proto-http")]
pub mod http;
#[cfg(feature = "proto-https")]
pub mod https;
#[cfg(feature = "proto-ssh")]
pub mod ssh;
#[cfg(feature = "hooks")]
pub mod hooks;
#[derive(Clone)]
pub struct DynamicProtocol {
pub name: String,
pub patterns: Vec<String>,
}
impl RefractiumProtocol for DynamicProtocol {
fn identify(self: Arc<Self>, data: &[u8]) -> Option<ProtocolMatch> {
use heck::ToSnakeCase;
let matched = self
.patterns
.iter()
.any(|p| memmem::find(data, p.as_bytes()).is_some());
if matched {
return Some(ProtocolMatch {
name: self.name.to_snake_case(),
metadata: None,
implementation: self,
});
}
None
}
fn name(&self) -> String {
self.name.clone()
}
fn transport(&self) -> Transport {
Transport::Both
}
}
pub struct ProtocolRegistry {
protocols: Vec<Arc<dyn RefractiumProtocol>>,
}
impl Default for ProtocolRegistry {
fn default() -> Self {
Self::new()
}
}
impl ProtocolRegistry {
#[must_use]
pub fn new() -> Self {
Self {
protocols: Vec::new(),
}
}
pub fn register(&mut self, proto: Arc<dyn RefractiumProtocol>) {
self.protocols.push(proto);
}
#[must_use]
pub fn get_registered_names(&self) -> Vec<String> {
use heck::ToSnakeCase;
self.protocols
.iter()
.map(|p| p.name().to_snake_case())
.collect()
}
#[must_use]
pub fn probe(&self, data: &[u8]) -> Option<ProtocolMatch> {
for proto in &self.protocols {
if let Some(m) = Arc::clone(proto).identify(data) {
return Some(m);
}
}
None
}
}