wifi-defibrillator.py

#!/usr/bin/env python3

# Copyright (C) 2023 Patrick Pedersen, TUDO Makerspace

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

# Author: Patrick Pedersen <[email protected]>
# Description:
#       The following script is based on the WiFi watchdog implemented in the TU-DO Makerspace
#       Activity Indicator, and is used to monitor the WiFi connection and reset the WiFi if
#       the connection is lost and could not be re-established after a timeout.

import os
import sys
import time

from enum import Enum
from ping3 import ping

WLAN_INTERFACE = "wlan0"
RESET_WIFI_TIMEOUT = 60

# Wrapper for print to work with journalctl logs
# -
# msg - Message to be printed
def print_journalctl(msg):
	print(msg)
	sys.stdout.flush() # Needed, else print won't appear in journalctl log

# Checks if internet connection is available
# -
# Returns True if connection is available, False otherwise
def check_connection():
	ret = ping('google.com')
	return ret != None and ret != False

# Waits until a connection to the internet is available
# -
# If a connection is available, true is returned.
# If a connection is not available after the provided timeout
# in seconds, false is returned.
# -
# timeout: timeout in seconds
# -
# Returns True if a connection is available, False if a connection is not available after the timeout
# NOTE: This function will block until a connection is available
def wait_for_connection(timeout_s: int):
	timeout = time.time() + timeout_s
	while not check_connection():
		if time.time() >= timeout:
			return False
	return True

# Handles WiFi connection related tasks, such as
# - Checking if a connection is available
# - Setting the connection indicator LED
# - Attempting to reconnect to the WiFi network
#
# If the functon fails to connect to the WiFi network after 60 seconds, it will
# raise a ConnectionError exception.
#
# NOTE: This function will block until a connection is available
def handle_connection():
	if not check_connection():
		if not wait_for_connection(RESET_WIFI_TIMEOUT):
			raise ConnectionError("No connection available after timeout")

# Resets the WiFi connection by restarting the WiFi interface
def reset_wifi():
	os.system("ifconfig wlan0 down")
	os.system("ifconfig wlan0 up")

prev_wifi_state = check_connection()

# Main loop
while True:
	try:
		handle_connection()
	except ConnectionError as e:
		print_journalctl("Failed to establish connection, resetting WiFi interface and retrying...")
		reset_wifi()
		prev_wifi_state = False
		continue

	if not prev_wifi_state:
		print_journalctl("Successfully re-established connection!")
		prev_wifi_state = True