first commit

This commit is contained in:
aOK 2024-06-14 14:40:06 +03:00
commit 827124d502
10 changed files with 3348 additions and 0 deletions

2526
Cargo.lock generated Normal file

File diff suppressed because it is too large Load diff

39
Cargo.toml Normal file
View file

@ -0,0 +1,39 @@
[package]
name = "notifymylimits"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
tokio = { version = "1", features = ["full"] }
notify-rust = { version = "4.10.0"}
# heim = { version = "0.0.11", default-fatures = false, features = ["host", "sensors", "net"] }
# heim-host = { version = "0.0.11" }
# heim = { git = "https://github.com/andodeki/heim.git", rev = "5ed73cf", default-fatures = false, features = ["host", "sensors", "net"]}
heim = { git = "https://github.com/heim-rs/heim.git", branch = "master", default-fatures = false, features = ["host", "sensors", "net"]}
heim-host = { git = "https://github.com/heim-rs/heim.git", branch = "master"}
battery = "0.7.8"
libc = "0.2.152"
futures = "0.3.30"
pcap = "1.1.0"
chrono = "0.4.31"
objc = "0.2.7"
psutil = { version = "3.3.0", optional = true }
procfs = { version = "0.16.0", optional = true }
cocoa = "0.25.0"
core-graphics = "0.23.2"
core-foundation = "0.9.4"
[features]
default = []
linux = ["procfs"]
android = ["procfs"]
macos = ["psutil"]
[patch.crates-io]
heim = { git = "https://github.com/andodeki/heim.git", rev = "5ed73cf"}
heim-host = { git = "https://github.com/andodeki/heim.git", rev = "5ed73cf"}
# heim = { git = "https://github.com/heim-rs/heim.git", rev = "b292f15"}
# notify-rust = { git = "https://github.com/hoodie/notify-rust.git", rev = "996953c"}
notify-rust = { git = "https://github.com/andodeki/notify-rust.git", rev = "175bd7c"}

0
README.md Normal file
View file

View file

@ -0,0 +1,14 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>Label</key>
<string>com.example.notifymylimit</string>
<key>ProgramArguments</key>
<array>
<string>/usr/local/bin/notifymylimit</string>
</array>
<key>RunAtLoad</key>
<true/>
</dict>
</plist>

1
error.txt Normal file
View file

@ -0,0 +1 @@

4
extra.txt Normal file
View file

@ -0,0 +1,4 @@
cargo build --release
mv target/release/mbslimit /usr/local/bin
create /Library/LaunchDaemons/com.example.notifymylimit.plist
sudo launchctl load /Library/LaunchDaemons/com.example.notifymylimit.plist

181
src/batterylibd.rs Normal file
View file

@ -0,0 +1,181 @@
extern crate battery;
extern crate futures;
extern crate libc;
extern crate tokio;
use std::error::Error;
use std::ffi::CString;
use std::io::Read;
use std::mem::MaybeUninit;
use std::os::raw::{c_char, c_int, c_void};
use std::os::unix::io::AsRawFd;
use std::os::unix::io::RawFd;
use std::pin::Pin;
use std::result::Result;
use std::task::Context;
use std::task::Poll;
use tokio::io::unix::AsyncFd;
use tokio::io::Interest;
use tokio::io::ReadBuf;
use tokio::io::{AsyncRead, AsyncReadExt};
use futures::ready;
pub const KEY: &str = "com.apple.system.powersources.timeremaining";
struct NotifyFd {
fd: RawFd,
pub token: c_int,
}
impl NotifyFd {
pub fn new(key: &str) -> Result<Self, Box<dyn Error>> {
let mut token = MaybeUninit::<c_int>::uninit();
let mut nfd = MaybeUninit::<RawFd>::uninit();
unsafe {
let key = CString::new(key).unwrap();
let r = notify_register_file_descriptor(
key.as_ptr(),
nfd.as_mut_ptr(),
0,
token.as_mut_ptr(),
);
if r != 0 {
return Err("notify_register_file_descriptor failed".into());
}
}
let token = unsafe { token.assume_init() };
let nfd = unsafe { nfd.assume_init() };
return Ok(NotifyFd {
fd: nfd,
token: token,
});
}
}
impl Read for NotifyFd {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
unsafe {
let r = libc::read(self.fd, buf.as_mut_ptr() as *mut c_void, buf.len());
if r == -1 {
return Err(std::io::Error::last_os_error());
} else {
return Ok(r as usize);
}
}
}
}
impl Drop for NotifyFd {
fn drop(&mut self) {
unsafe {
let r = notify_cancel(self.token);
if r != 0 {
panic!("notify_cancel failed");
}
}
}
}
// Needed for integration with Tokio
impl AsRawFd for NotifyFd {
fn as_raw_fd(&self) -> RawFd {
return self.fd;
}
}
pub struct AsyncNotifyFd {
inner: AsyncFd<NotifyFd>,
pub token: c_int,
}
impl AsyncNotifyFd {
pub fn new(key: &str) -> Result<Self, Box<dyn Error>> {
let mut nfd = NotifyFd::new(key)?;
// Suspend the events while we adjust the fd
unsafe {
let r = notify_suspend(nfd.token);
if r != 0 {
return Err("notify_suspend failed".into());
}
}
// Set the file descriptor in non blocking mode
unsafe {
let flags = libc::fcntl(nfd.fd, libc::F_GETFL);
let r = libc::fcntl(nfd.fd, libc::F_SETFL, flags | libc::O_NONBLOCK);
if r != 0 {
return Err("fcntl failed".into());
}
}
// Drain the file descriptor of all data before registering with Tokio
loop {
let mut buf = [0; 4];
match nfd.read_exact(&mut buf) {
Ok(_) => {
continue;
}
Err(e) => {
if e.kind() == std::io::ErrorKind::WouldBlock {
break;
} else {
return Err(format!("unexpected read io error {}", e).into());
}
}
}
}
let t = nfd.token;
// Register the file descriptor with tokio
let afd = AsyncFd::with_interest(nfd, Interest::READABLE)?;
// Resume events
unsafe {
let r = notify_resume(t);
if r != 0 {
return Err("notify_resume failed".into());
}
}
return Ok(Self {
inner: afd,
token: t,
});
}
}
impl AsyncRead for AsyncNotifyFd {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
loop {
let mut guard = ready!(self.inner.poll_read_ready_mut(cx))?;
let r = guard.try_io(|x| x.get_mut().read(buf.initialize_unfilled()));
if r.is_ok() {
return Poll::Ready(r.unwrap().map(|r| buf.advance(r)));
}
}
}
}
extern "C" {
pub fn notify_register_file_descriptor(
name: *const c_char,
notify_fd: *mut c_int,
flags: c_int,
out_token: *mut c_int,
) -> u32;
pub fn notify_cancel(token: c_int) -> u32;
// Added to allow safely setting the fd to non blocking mode
pub fn notify_suspend(token: ::std::os::raw::c_int) -> u32;
pub fn notify_resume(token: ::std::os::raw::c_int) -> u32;
}

58
src/getforegroundapp.rs Normal file
View file

@ -0,0 +1,58 @@
use std::collections::HashSet;
use std::process::Command;
// Common function that delegates to the appropriate platform-specific implementation
pub fn get_foreground_process_ids() -> HashSet<i32> {
#[cfg(any(target_os = "linux", target_os = "windows"))]
return procfs_process_info::get_foreground_process_ids();
#[cfg(target_os = "macos")]
return macos_process_info::get_foreground_process_ids();
}
// Module containing platform-specific process information logic
#[cfg(any(target_os = "linux", target_os = "windows"))]
mod procfs_process_info {
use procfs::process::{FDInfo, FDTarget, Process};
// Function to get the set of foreground process IDs on Linux and Windows
pub fn get_foreground_process_ids() -> HashSet<i32> {
let terminal_fd_info = match Process::myself() {
Ok(process) => process.fd,
Err(_) => return HashSet::new(),
};
let mut foreground_process_ids = HashSet::new();
for fd_info in terminal_fd_info {
if let FDTarget::Process(process_info) = fd_info.target {
foreground_process_ids.insert(process_info.pid);
}
}
foreground_process_ids
}
}
#[cfg(target_os = "macos")]
mod macos_process_info {
use std::collections::HashSet;
use std::process::Command;
// Function to get the set of foreground process IDs on macOS
pub fn get_foreground_process_ids() -> HashSet<i32> {
let output = Command::new("sh")
.arg("-c")
.arg("lsappinfo info -only pid $(lsappinfo front)")
.output()
.expect("Failed to execute command");
let pid: i32 = String::from_utf8(output.stdout)
.unwrap()
.trim()
.parse()
.unwrap();
let mut foreground_process_ids = HashSet::new();
foreground_process_ids.insert(pid);
foreground_process_ids
}
}

2
src/lib.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod batterylibd;
pub mod getforegroundapp;

523
src/main.rs Normal file
View file

@ -0,0 +1,523 @@
use cocoa::{
base::{id, nil},
foundation::{NSAutoreleasePool, NSString},
};
use notify_rust::Notification;
use objc::{
declare::ClassDecl,
msg_send,
runtime::{Class, Object, Sel},
sel, sel_impl,
};
use pcap::Capture;
use std::{
collections::HashMap,
fs::{File, OpenOptions},
io::{self, Write},
process::Command,
sync::{Arc, Mutex},
time::Duration,
};
use tokio::{io::unix::AsyncFd, time::sleep};
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
// if let Err(mbs_error) = mbslimit().await {
// eprintln!("MBS error: {:?}", mbs_error);
// }
//
if let Err(foreground_app_error) = monitor_foreground_app_changes().await {
eprintln!("Foreground App error: {:?}", foreground_app_error);
}
// let (mbs_result, battery_result, monitor_result) =
// tokio::join!(mbslimit(), batterylimit(), monitor_foreground_app_changes());
// if let Err(e) = mbs_result {
// eprintln!("Error in mbslimit: {:?}", e);
// }
// if let Err(e) = battery_result {
// eprintln!("Error in batterylimit: {:?}", e);
// }
// if let Err(e) = monitor_result {
// eprintln!("Error in monitor_foreground_app_changes: {:?}", e);
// }
Ok(())
}
fn get_foreground_app() -> Option<String> {
unsafe {
let pool = NSAutoreleasePool::new(nil);
let workspace: id = msg_send![Class::get("NSWorkspace").unwrap(), sharedWorkspace];
let app: id = msg_send![workspace, frontmostApplication];
let bundle_id: id = msg_send![app, bundleIdentifier];
if bundle_id != nil {
let bundle_id_str: *const std::os::raw::c_char = msg_send![bundle_id, UTF8String];
let app_name = std::ffi::CStr::from_ptr(bundle_id_str)
.to_string_lossy()
.into_owned();
pool.drain();
Some(app_name)
} else {
pool.drain();
None
}
}
}
fn block_app(app_bundle_id: &str) -> Result<(), Box<dyn std::error::Error>> {
let rule = format!(
"block out quick proto tcp from any to any app \"{}\"",
app_bundle_id
);
Command::new("sudo")
.arg("pfctl")
.arg("-a")
.arg("foreground_app")
.arg("-f")
.arg(rule)
.output()?;
Command::new("sudo").arg("pfctl").arg("-e").output()?;
Ok(())
}
fn log_data_usage(app_bundle_id: &str, data_usage: usize) -> io::Result<()> {
let mut file = OpenOptions::new()
.create(true)
.append(true)
.open("data_usage_history.txt")?;
writeln!(file, "{}: {} bytes", app_bundle_id, data_usage)?;
Ok(())
}
async fn monitor_foreground_app_changes() -> Result<(), Box<dyn std::error::Error>> {
let mut current_app = get_foreground_app().unwrap_or_else(|| "Unknown".to_string());
let mut app_data_usage: HashMap<String, usize> = HashMap::new();
app_data_usage.insert(current_app.clone(), 0);
loop {
if let Some(new_app) = get_foreground_app() {
if new_app != current_app {
println!("Foreground app changed from {} to {}", current_app, new_app);
if let Some(&usage) = app_data_usage.get(&current_app) {
log_data_usage(&current_app, usage)?;
println!("{} used {} bytes", current_app, usage);
}
block_app(&current_app)?;
current_app = new_app;
app_data_usage.entry(current_app.clone()).or_insert(0);
}
}
sleep(Duration::from_secs(1)).await;
}
}
async fn mbslimit() -> Result<(), Box<dyn std::error::Error>> {
println!("Entering mbslimit function");
let device = pcap::Device::lookup()?.ok_or("No device available")?;
println!("Using device {}", device.name);
let mut cap = Capture::from_device(device)?.immediate_mode(true).open()?;
let mut total_size = 0;
let mut current_app = get_foreground_app().unwrap_or_else(|| "Unknown".to_string());
let mut app_data_usage: HashMap<String, usize> = HashMap::new();
app_data_usage.insert(current_app.clone(), 0);
loop {
while let Ok(packet) = cap.next_packet() {
total_size += packet.len();
*app_data_usage.entry(current_app.clone()).or_insert(0) += packet.len();
if total_size > 24 * 1024 * 1024 {
total_size = 0;
println!("Alert: Total data transferred exceeded 24MB");
notify_and_toggle_wifi().await?;
break;
}
}
if let Some(new_app) = get_foreground_app() {
if new_app != current_app {
println!("Foreground app changed from {} to {}", current_app, new_app);
if let Some(&usage) = app_data_usage.get(&current_app) {
log_data_usage(&current_app, usage)?;
println!("{} used {} bytes", current_app, usage);
}
block_app(&current_app)?;
current_app = new_app;
app_data_usage.entry(current_app.clone()).or_insert(0);
}
}
sleep(Duration::from_secs(1)).await;
}
}
async fn notify_and_toggle_wifi() -> Result<(), Box<dyn std::error::Error>> {
Command::new("osascript")
.arg("-e")
.arg("display notification \"Data transfer limit exceeded\" with title \"Alert\"")
.output()?;
Command::new("networksetup")
.arg("-setairportpower")
.arg("en0")
.arg("off")
.output()?;
Notification::new()
.summary("WiFi Off")
.body("Data usage has exceeded 24MB. WiFi has been turned off.")
.show()?;
println!("WiFi has been turned off.");
sleep(Duration::from_secs(5 * 60)).await;
Command::new("networksetup")
.arg("-setairportpower")
.arg("en0")
.arg("on")
.output()?;
Notification::new()
.summary("WiFi On")
.body("WiFi has been turned back on.")
.show()?;
println!("WiFi has been turned on.");
Ok(())
}
async fn batterylimit() -> Result<(), Box<dyn std::error::Error>> {
println!("Entering batterylimit function");
let manager = battery::Manager::new()?;
let mut battery = manager.batteries()?.next().ok_or("No battery found")??;
let mut nfd = AsyncNotifyFd::new(KEY)?;
let mut buf = [0; 4];
loop {
nfd.read_exact(&mut buf).await?;
let v = c_int::from_be_bytes(buf);
if v == nfd.token {
manager
.refresh(&mut battery)
.expect("Manager failed to refresh");
} else {
return Err("Unknown token in file descriptor!".into());
}
let percentage = battery.state_of_charge().value * 100.0;
if percentage > 70.0 {
notify_battery_level(percentage, "high").await?;
} else if percentage < 25.0 {
notify_battery_level(percentage, "low").await?;
}
sleep(Duration::from_secs(1)).await;
}
}
async fn notify_battery_level(
percentage: f32,
level: &str,
) -> Result<(), Box<dyn std::error::Error>> {
Notification::new()
.summary("Battery Level Alert")
.body(&format!(
"Your battery level is now at {:.0}%. Please monitor it.",
percentage
))
.show()?;
println!("Battery charge: {:.0}%", percentage);
Ok(())
}
fn block_all_except_foreground_app(app_bundle_id: &str) -> Result<(), Box<dyn std::error::Error>> {
// Load a custom pf.conf file that blocks all traffic by default
// pfctl -f /etc/pf.conf
Command::new("sudo")
.arg("pfctl")
.arg("-f")
.arg("/etc/pf.conf")
.output()?;
// Allow traffic for the foreground app (you might need to tailor this to your exact needs)
let rule = format!(
"pass out quick proto tcp from any to any app \"{}\" keep state",
app_bundle_id
);
// sudo pfctl -a foreground_app -f
Command::new("sudo")
.arg("pfctl")
.arg("-a")
.arg("foreground_app")
.arg("-f")
.arg(rule)
.output()?;
// Reload PF configuration
Command::new("sudo").arg("pfctl").arg("-e").output()?;
Ok(())
}
// fn get_foreground_app() -> Option<String> {
// unsafe {
// let pool = NSAutoreleasePool::new(nil);
// let workspace: id = msg_send![Class::get("NSWorkspace").unwrap(), sharedWorkspace];
// let app: id = msg_send![workspace, frontmostApplication];
// let bundle_id: id = msg_send![app, bundleIdentifier];
// if bundle_id != nil {
// let bundle_id_str: *const std::os::raw::c_char = msg_send![bundle_id, UTF8String];
// let app_name = std::ffi::CStr::from_ptr(bundle_id_str)
// .to_string_lossy()
// .into_owned();
// pool.drain();
// Some(app_name)
// } else {
// pool.drain();
// None
// }
// }
// }
// fn block_app(app_bundle_id: &str) -> Result<(), Box<dyn std::error::Error>> {
// // Block traffic for the app (you might need to tailor this to your exact needs)
// let rule = format!(
// "block out quick proto tcp from any to any app \"{}\"",
// app_bundle_id
// );
// Command::new("sudo")
// .arg("pfctl")
// .arg("-a")
// .arg("foreground_app")
// .arg("-f")
// .arg(rule)
// .output()?;
// // Reload PF configuration
// Command::new("sudo").arg("pfctl").arg("-e").output()?;
// Ok(())
// }
// fn log_data_usage(app_bundle_id: &str, data_usage: usize) -> io::Result<()> {
// let mut file = OpenOptions::new()
// .create(true)
// .append(true)
// .open("data_usage_history.txt")?;
// writeln!(file, "{}: {} bytes", app_bundle_id, data_usage)?;
// Ok(())
// }
// async fn notify_and_toggle_wifi() -> Result<(), Box<dyn std::error::Error>> {
// // Notify user
// Command::new("osascript")
// .arg("-e")
// .arg("display notification \"Data transfer limit exceeded\" with title \"Alert\"")
// .output()?;
// // Turn off WiFi
// Command::new("networksetup")
// .arg("-setairportpower")
// .arg("en0")
// .arg("off")
// .output()?;
// Notification::new()
// .summary("WiFi Off")
// .body("Data usage has exceeded 24MB. WiFi has been turned off.")
// .show()?;
// // Wait for 5 minutes
// println!("WiFi has been turned off.");
// sleep(Duration::from_secs(5 * 60)).await;
// // Turn on WiFi
// Command::new("networksetup")
// .arg("-setairportpower")
// .arg("en0")
// .arg("on")
// .output()?;
// Notification::new()
// .summary("WiFi On")
// .body("WiFi has been turned back on.")
// .show()?;
// println!("WiFi has been turned on.");
// Ok(())
// }
extern crate battery;
extern crate futures;
extern crate libc;
extern crate tokio;
use futures::ready;
use std::error::Error;
use std::ffi::CString;
use std::io::Read;
use std::mem::MaybeUninit;
use std::os::raw::{c_char, c_int, c_void};
use std::os::unix::io::AsRawFd;
use std::os::unix::io::RawFd;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncReadExt, Interest, ReadBuf};
pub const KEY: &str = "com.apple.system.powersources.timeremaining";
struct NotifyFd {
fd: RawFd,
token: c_int,
}
impl NotifyFd {
pub fn new(key: &str) -> Result<Self, Box<dyn Error>> {
let mut token = MaybeUninit::<c_int>::uninit();
let mut nfd = MaybeUninit::<RawFd>::uninit();
unsafe {
let key = CString::new(key)?;
let r = notify_register_file_descriptor(
key.as_ptr(),
nfd.as_mut_ptr(),
0,
token.as_mut_ptr(),
);
if r != 0 {
return Err("notify_register_file_descriptor failed".into());
}
}
Ok(NotifyFd {
fd: unsafe { nfd.assume_init() },
token: unsafe { token.assume_init() },
})
}
}
impl Read for NotifyFd {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let r = unsafe { libc::read(self.fd, buf.as_mut_ptr() as *mut c_void, buf.len()) };
if r == -1 {
Err(std::io::Error::last_os_error())
} else {
Ok(r as usize)
}
}
}
impl Drop for NotifyFd {
fn drop(&mut self) {
unsafe {
let r = notify_cancel(self.token);
if r != 0 {
panic!("notify_cancel failed");
}
}
}
}
impl AsRawFd for NotifyFd {
fn as_raw_fd(&self) -> RawFd {
self.fd
}
}
pub struct AsyncNotifyFd {
inner: AsyncFd<NotifyFd>,
token: c_int,
}
impl AsyncNotifyFd {
pub fn new(key: &str) -> Result<Self, Box<dyn Error>> {
let mut nfd = NotifyFd::new(key)?;
unsafe {
let r = notify_suspend(nfd.token);
if r != 0 {
return Err("notify_suspend failed".into());
}
let flags = libc::fcntl(nfd.fd, libc::F_GETFL);
let r = libc::fcntl(nfd.fd, libc::F_SETFL, flags | libc::O_NONBLOCK);
if r != 0 {
return Err("fcntl failed".into());
}
}
loop {
let mut buf = [0; 4];
match nfd.read_exact(&mut buf) {
Ok(_) => continue,
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
Err(e) => return Err(format!("unexpected read io error {}", e).into()),
}
}
let t = nfd.token;
let afd = AsyncFd::with_interest(nfd, Interest::READABLE)?;
unsafe {
let r = notify_resume(t);
if r != 0 {
return Err("notify_resume failed".into());
}
}
Ok(Self {
inner: afd,
token: t,
})
}
}
impl AsyncRead for AsyncNotifyFd {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
loop {
let mut guard = ready!(self.inner.poll_read_ready_mut(cx))?;
let r = guard.try_io(|inner| inner.get_mut().read(buf.initialize_unfilled()));
if r.is_ok() {
return Poll::Ready(r.unwrap().map(|n| buf.advance(n)));
}
}
}
}
extern "C" {
fn notify_register_file_descriptor(
name: *const c_char,
notify_fd: *mut c_int,
flags: c_int,
out_token: *mut c_int,
) -> u32;
fn notify_cancel(token: c_int) -> u32;
fn notify_suspend(token: c_int) -> u32;
fn notify_resume(token: c_int) -> u32;
}