#!/usr/bin/env python3
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import asyncio
import psutil
import time
import os
import argparse
import logging
from datetime import datetime
from kasa import SmartStrip
class KasaPlugController():
"""Provides control of a device's charger.
The device's charger must be plugged into one of the 3 outlets of a Kasa Smart
Plug Power Strip (KP303). The outlet name must match the device's host name
(this is intended to prevent inadvertently controlling the wrong device's
charger).
"""
def __init__(self, kasa_power_strip_ip: str):
"""Constructs a KasaPlugController to control the current device's charger.
Args:
kasa_power_strip_ip: IP of the Kasak Smart Plug Power Strip in which
this device's charger is connected.
"""
# The outlet name must match the device's host name.
self.kasa_outlet_name = os.uname()[1].split('.')[0]
# Create the event loop
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)
# Create the strip controller
self.strip = SmartStrip(kasa_power_strip_ip)
self.loop.run_until_complete(self.strip.update())
self.closed = False
def __del__(self):
self.close()
def turn_on(self):
"""Turns on this device's charger.
"""
logging.info("Turning on the charger")
for plug in self.strip.children:
if plug.alias == self.kasa_outlet_name:
self.loop.run_until_complete(plug.turn_on())
return
logging.error("Cannot find device to turn on")
def turn_off(self):
"""Turns off this device's charger.
"""
logging.info("Turning off the charger")
for plug in self.strip.children:
if plug.alias == self.kasa_outlet_name:
self.loop.run_until_complete(plug.turn_off())
battery = psutil.sensors_battery()
while battery.power_plugged:
logging.info("Waiting for device to no longer be plugged in")
time.sleep(1)
battery = psutil.sensors_battery()
return
logging.error("Cannot find device to turn off")
def discharge_to(self, level: int):
"""Discharges the battery until it reaches a target level.
Args:
level: The target battery level.
"""
self.turn_off()
battery = psutil.sensors_battery()
while battery.percent > level:
logging.info(f"Waiting to discharge to {level}%."
f" Currently at {battery.percent}%")
# Perform arbitrary operations as fast as possible to burn
# CPU and discharge faster.
f_value = 0.81
start = datetime.now()
while ((datetime.now() - start).total_seconds() < 10):
f_value = f_value * 1.7272882
f_value = f_value / 1.7272882
battery = psutil.sensors_battery()
logging.info(f"Discharge to {level}% complete")
def charge_to(self, level: int):
"""Charges the battery until it reaches a target level.
Args:
level: The target battery level.
"""
self.turn_on()
battery = psutil.sensors_battery()
while battery.percent < level:
logging.info(f"Waiting to charge to {level}%."
f" Currently at {battery.percent}%")
time.sleep(10)
battery = psutil.sensors_battery()
logging.info(f"Charge to {level}% complete")
def charge_or_discharge_to(self, level: int):
"""Charges or discharges the battery until it reaches a target level.
Leaves the charger in an unplugged state.
Args:
level: The target battery level.
"""
battery = psutil.sensors_battery()
if battery.percent < level:
self.charge_to(level)
elif battery.percent > level:
self.discharge_to(level)
else:
logging.info(f"Battery is already at the target level {level}%")
self.turn_off()
def close(self):
"""Closes the message loop."""
if self.closed:
return
self.closed = True
self.loop.close()
def get_plug_controller(ip: str):
return KasaPlugController(ip)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description='Controls kasa power switch connected to this device.')
parser.add_argument("--kasa_power_strip_ip",
required=True,
help="IP address of the kasa power switch.")
parser.add_argument("--charge_level",
type=int,
required=True,
help="Desired charge level.")
args = parser.parse_args()
kasa_plug_controller = KasaPlugController(args.kasa_power_strip_ip)
kasa_plug_controller.charge_to(args.charge_level)
kasa_plug_controller.close()