Skip to content

Commit

Permalink
[feature] Structured Reporting - Files (#496)
Browse files Browse the repository at this point in the history
* update protobufs

* Use c2pb enum on process ent, rename to HostProcess, add fields

* Migrate to c2pb for Host_Platform enum

* Fix enums, added tests

* Added ReportFile gRPC API

* Generate Rust Bindings

* minor cleanup
  • Loading branch information
KCarretto authored Jan 23, 2024
1 parent 3d92ba2 commit 7fd5b31
Show file tree
Hide file tree
Showing 107 changed files with 23,305 additions and 7,921 deletions.
259 changes: 201 additions & 58 deletions implants/lib/c2/src/c2.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,26 @@
/// Agent information to identify the type of beacon.
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct Agent {
#[prost(string, tag = "1")]
pub identifier: ::prost::alloc::string::String,
}
/// Beacon information that is unique to the current running beacon.
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct Beacon {
#[prost(string, tag = "1")]
pub identifier: ::prost::alloc::string::String,
#[prost(string, tag = "2")]
pub principal: ::prost::alloc::string::String,
#[prost(message, optional, tag = "3")]
pub host: ::core::option::Option<Host>,
#[prost(message, optional, tag = "4")]
pub agent: ::core::option::Option<Agent>,
/// Duration until next callback, in seconds.
#[prost(uint64, tag = "5")]
pub interval: u64,
}
/// Host information for the system a beacon is running on.
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
Expand Down Expand Up @@ -59,28 +82,103 @@ pub mod host {
}
}
}
/// Agent information to identify the type of beacon.
/// Process running on the host system.
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct Agent {
#[prost(string, tag = "1")]
pub identifier: ::prost::alloc::string::String,
}
/// Beacon information that is unique to the current running beacon.
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct Beacon {
#[prost(string, tag = "1")]
pub identifier: ::prost::alloc::string::String,
#[prost(string, tag = "2")]
pub struct Process {
#[prost(uint64, tag = "1")]
pub pid: u64,
#[prost(uint64, tag = "2")]
pub ppid: u64,
#[prost(string, tag = "3")]
pub name: ::prost::alloc::string::String,
#[prost(string, tag = "4")]
pub principal: ::prost::alloc::string::String,
#[prost(message, optional, tag = "3")]
pub host: ::core::option::Option<Host>,
#[prost(message, optional, tag = "4")]
pub agent: ::core::option::Option<Agent>,
/// Duration until next callback, in seconds.
#[prost(uint64, tag = "5")]
pub interval: u64,
#[prost(string, tag = "5")]
pub path: ::prost::alloc::string::String,
#[prost(string, tag = "6")]
pub cmd: ::prost::alloc::string::String,
#[prost(string, tag = "7")]
pub env: ::prost::alloc::string::String,
#[prost(string, tag = "8")]
pub cwd: ::prost::alloc::string::String,
#[prost(enumeration = "process::Status", tag = "9")]
pub status: i32,
}
/// Nested message and enum types in `Process`.
pub mod process {
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
Hash,
PartialOrd,
Ord,
::prost::Enumeration
)]
#[repr(i32)]
pub enum Status {
Unspecified = 0,
Unknown = 1,
Idle = 2,
Run = 3,
Sleep = 4,
Stop = 5,
Zombie = 6,
Tracing = 7,
Dead = 8,
WakeKill = 9,
Waking = 10,
Parked = 11,
LockBlocked = 12,
UninteruptibleDiskSleep = 13,
}
impl Status {
/// String value of the enum field names used in the ProtoBuf definition.
///
/// The values are not transformed in any way and thus are considered stable
/// (if the ProtoBuf definition does not change) and safe for programmatic use.
pub fn as_str_name(&self) -> &'static str {
match self {
Status::Unspecified => "STATUS_UNSPECIFIED",
Status::Unknown => "STATUS_UNKNOWN",
Status::Idle => "STATUS_IDLE",
Status::Run => "STATUS_RUN",
Status::Sleep => "STATUS_SLEEP",
Status::Stop => "STATUS_STOP",
Status::Zombie => "STATUS_ZOMBIE",
Status::Tracing => "STATUS_TRACING",
Status::Dead => "STATUS_DEAD",
Status::WakeKill => "STATUS_WAKE_KILL",
Status::Waking => "STATUS_WAKING",
Status::Parked => "STATUS_PARKED",
Status::LockBlocked => "STATUS_LOCK_BLOCKED",
Status::UninteruptibleDiskSleep => "STATUS_UNINTERUPTIBLE_DISK_SLEEP",
}
}
/// Creates an enum from field names used in the ProtoBuf definition.
pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
match value {
"STATUS_UNSPECIFIED" => Some(Self::Unspecified),
"STATUS_UNKNOWN" => Some(Self::Unknown),
"STATUS_IDLE" => Some(Self::Idle),
"STATUS_RUN" => Some(Self::Run),
"STATUS_SLEEP" => Some(Self::Sleep),
"STATUS_STOP" => Some(Self::Stop),
"STATUS_ZOMBIE" => Some(Self::Zombie),
"STATUS_TRACING" => Some(Self::Tracing),
"STATUS_DEAD" => Some(Self::Dead),
"STATUS_WAKE_KILL" => Some(Self::WakeKill),
"STATUS_WAKING" => Some(Self::Waking),
"STATUS_PARKED" => Some(Self::Parked),
"STATUS_LOCK_BLOCKED" => Some(Self::LockBlocked),
"STATUS_UNINTERUPTIBLE_DISK_SLEEP" => Some(Self::UninteruptibleDiskSleep),
_ => None,
}
}
}
}
/// Task instructions for the beacon to execute.
#[allow(clippy::derive_partial_eq_without_eq)]
Expand Down Expand Up @@ -124,17 +222,6 @@ pub struct TaskOutput {
#[prost(message, optional, tag = "5")]
pub exec_finished_at: ::core::option::Option<::prost_types::Timestamp>,
}
/// Process running on the system.
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct Process {
#[prost(uint64, tag = "1")]
pub pid: u64,
#[prost(string, tag = "2")]
pub name: ::prost::alloc::string::String,
#[prost(string, tag = "3")]
pub principal: ::prost::alloc::string::String,
}
///
/// RPC Messages
#[allow(clippy::derive_partial_eq_without_eq)]
Expand All @@ -151,15 +238,6 @@ pub struct ClaimTasksResponse {
}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportTaskOutputRequest {
#[prost(message, optional, tag = "1")]
pub output: ::core::option::Option<TaskOutput>,
}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportTaskOutputResponse {}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct DownloadFileRequest {
#[prost(string, tag = "1")]
pub name: ::prost::alloc::string::String,
Expand All @@ -172,6 +250,29 @@ pub struct DownloadFileResponse {
}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportFileRequest {
#[prost(int64, tag = "1")]
pub task_id: i64,
#[prost(string, tag = "2")]
pub path: ::prost::alloc::string::String,
#[prost(string, tag = "3")]
pub owner: ::prost::alloc::string::String,
#[prost(string, tag = "4")]
pub group: ::prost::alloc::string::String,
#[prost(string, tag = "5")]
pub permissions: ::prost::alloc::string::String,
#[prost(int64, tag = "6")]
pub size: i64,
#[prost(string, tag = "7")]
pub sha3_256_hash: ::prost::alloc::string::String,
#[prost(bytes = "vec", tag = "8")]
pub chunk: ::prost::alloc::vec::Vec<u8>,
}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportFileResponse {}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportProcessListRequest {
#[prost(message, repeated, tag = "1")]
pub list: ::prost::alloc::vec::Vec<Process>,
Expand All @@ -181,6 +282,15 @@ pub struct ReportProcessListRequest {
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportProcessListResponse {}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportTaskOutputRequest {
#[prost(message, optional, tag = "1")]
pub output: ::core::option::Option<TaskOutput>,
}
#[allow(clippy::derive_partial_eq_without_eq)]
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ReportTaskOutputResponse {}
/// Generated client implementations.
pub mod c2_client {
#![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
Expand Down Expand Up @@ -266,6 +376,8 @@ pub mod c2_client {
self.inner = self.inner.max_encoding_message_size(limit);
self
}
///
/// Contact the server for new tasks to execute.
pub async fn claim_tasks(
&mut self,
request: impl tonic::IntoRequest<super::ClaimTasksRequest>,
Expand All @@ -288,11 +400,19 @@ pub mod c2_client {
req.extensions_mut().insert(GrpcMethod::new("c2.C2", "ClaimTasks"));
self.inner.unary(req, path, codec).await
}
pub async fn report_task_output(
///
/// Download a file from the server, returning one or more chunks of data.
/// The maximum size of these chunks is determined by the server.
/// The server should reply with two headers:
/// - "sha3-256-checksum": A SHA3-256 digest of the entire file contents.
/// - "file-size": The number of bytes contained by the file.
///
/// If no associated file can be found, a NotFound status error is returned.
pub async fn download_file(
&mut self,
request: impl tonic::IntoRequest<super::ReportTaskOutputRequest>,
request: impl tonic::IntoRequest<super::DownloadFileRequest>,
) -> std::result::Result<
tonic::Response<super::ReportTaskOutputResponse>,
tonic::Response<tonic::codec::Streaming<super::DownloadFileResponse>>,
tonic::Status,
> {
self.inner
Expand All @@ -305,10 +425,39 @@ pub mod c2_client {
)
})?;
let codec = tonic::codec::ProstCodec::default();
let path = http::uri::PathAndQuery::from_static("/c2.C2/ReportTaskOutput");
let path = http::uri::PathAndQuery::from_static("/c2.C2/DownloadFile");
let mut req = request.into_request();
req.extensions_mut().insert(GrpcMethod::new("c2.C2", "ReportTaskOutput"));
self.inner.unary(req, path, codec).await
req.extensions_mut().insert(GrpcMethod::new("c2.C2", "DownloadFile"));
self.inner.server_streaming(req, path, codec).await
}
///
/// Report a file from the host to the server.
/// Providing content of the file is optional. If content is provided:
/// - Hash will automatically be calculated and the provided hash will be ignored.
/// - Size will automatically be calculated and the provided size will be ignored.
/// Content is provided as chunks, the size of which are up to the agent to define (based on memory constraints).
/// Any existing files at the provided path for the host are replaced.
pub async fn report_file(
&mut self,
request: impl tonic::IntoStreamingRequest<Message = super::ReportFileRequest>,
) -> std::result::Result<
tonic::Response<super::ReportFileResponse>,
tonic::Status,
> {
self.inner
.ready()
.await
.map_err(|e| {
tonic::Status::new(
tonic::Code::Unknown,
format!("Service was not ready: {}", e.into()),
)
})?;
let codec = tonic::codec::ProstCodec::default();
let path = http::uri::PathAndQuery::from_static("/c2.C2/ReportFile");
let mut req = request.into_streaming_request();
req.extensions_mut().insert(GrpcMethod::new("c2.C2", "ReportFile"));
self.inner.client_streaming(req, path, codec).await
}
///
/// Report the active list of running processes. This list will replace any previously reported
Expand Down Expand Up @@ -336,18 +485,12 @@ pub mod c2_client {
self.inner.unary(req, path, codec).await
}
///
/// Download a file from the server, returning one or more chunks of data.
/// The maximum size of these chunks is determined by the server.
/// The server should reply with two headers:
/// - "sha3-256-checksum": A SHA3-256 digest of the entire file contents.
/// - "file-size": The number of bytes contained by the file.
///
/// If no associated file can be found, a NotFound status error is returned.
pub async fn download_file(
/// Report execution output for a task.
pub async fn report_task_output(
&mut self,
request: impl tonic::IntoRequest<super::DownloadFileRequest>,
request: impl tonic::IntoRequest<super::ReportTaskOutputRequest>,
) -> std::result::Result<
tonic::Response<tonic::codec::Streaming<super::DownloadFileResponse>>,
tonic::Response<super::ReportTaskOutputResponse>,
tonic::Status,
> {
self.inner
Expand All @@ -360,10 +503,10 @@ pub mod c2_client {
)
})?;
let codec = tonic::codec::ProstCodec::default();
let path = http::uri::PathAndQuery::from_static("/c2.C2/DownloadFile");
let path = http::uri::PathAndQuery::from_static("/c2.C2/ReportTaskOutput");
let mut req = request.into_request();
req.extensions_mut().insert(GrpcMethod::new("c2.C2", "DownloadFile"));
self.inner.server_streaming(req, path, codec).await
req.extensions_mut().insert(GrpcMethod::new("c2.C2", "ReportTaskOutput"));
self.inner.unary(req, path, codec).await
}
}
}
3 changes: 1 addition & 2 deletions tavern/internal/c2/api_claim_tasks.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,12 @@ func (srv *Server) ClaimTasks(ctx context.Context, req *c2pb.ClaimTasksRequest)
if req.Beacon.Agent.Identifier == "" {
return nil, status.Errorf(codes.InvalidArgument, "must provide agent identifier")
}
hostPlaform := convertHostPlatform(req.Beacon.Host.Platform)

// Upsert the host
hostID, err := srv.graph.Host.Create().
SetIdentifier(req.Beacon.Host.Identifier).
SetName(req.Beacon.Host.Name).
SetPlatform(hostPlaform).
SetPlatform(req.Beacon.Host.Platform).
SetPrimaryIP(req.Beacon.Host.PrimaryIp).
SetLastSeenAt(now).
OnConflict().
Expand Down
Loading

0 comments on commit 7fd5b31

Please sign in to comment.