switched to Arduino
This commit is contained in:
parent
b69f3d8b2a
commit
a82c96de1a
19
Makefile
19
Makefile
@ -5,24 +5,7 @@ SERIAL_PORT=/dev/tty.usbserial-1410
|
||||
#SERIAL_PORT=/dev/tty.SLAB_USBtoUART
|
||||
#SERIAL_PORT=/dev/tty.wchusbserial1410
|
||||
|
||||
all: flash delay libs config update reset
|
||||
|
||||
delay:
|
||||
sleep 3
|
||||
|
||||
reset:
|
||||
esptool.py -p $(SERIAL_PORT) --after hard_reset read_mac
|
||||
|
||||
libs:
|
||||
ampy -d 0.5 -p $(SERIAL_PORT) put uwebsockets.py
|
||||
|
||||
update:
|
||||
ampy -d 0.5 -p $(SERIAL_PORT) put hal.py
|
||||
ampy -d 0.5 -p $(SERIAL_PORT) put main.py
|
||||
ampy -d 0.5 -p $(SERIAL_PORT) put boot.py
|
||||
|
||||
config:
|
||||
ampy -d 0.5 -p $(SERIAL_PORT) put config.json
|
||||
all: flash
|
||||
|
||||
flash:
|
||||
esptool.py -p $(SERIAL_PORT) -b 460800 erase_flash
|
||||
|
53
boot.py
53
boot.py
@ -1,53 +0,0 @@
|
||||
# Import functions from time library
|
||||
from utime import ticks_us, sleep_us, sleep_ms
|
||||
|
||||
# Give time to cancel boot script
|
||||
print("Press Ctrl-C to stop boot script...")
|
||||
sleep_ms(200)
|
||||
|
||||
# Import libraries
|
||||
import os
|
||||
import ujson
|
||||
import network
|
||||
import _thread
|
||||
import uwebsockets
|
||||
from machine import Timer, reset
|
||||
from hal import *
|
||||
# Loading libraries takes ca 400ms
|
||||
|
||||
# Open and parse the config file
|
||||
with open("config.json", "r") as config_file:
|
||||
config = ujson.load(config_file)
|
||||
|
||||
# Initialize the SumoRobot object
|
||||
sumorobot = Sumorobot(config)
|
||||
|
||||
# Indiacte booting with blinking status LED
|
||||
timer = Timer(0)
|
||||
timer.init(period=2000, mode=Timer.PERIODIC, callback=sumorobot.toggle_led)
|
||||
|
||||
# Connected Wi-Fi SSID
|
||||
ssid = None
|
||||
|
||||
# Connect to WiFi
|
||||
wlan = network.WLAN(network.STA_IF)
|
||||
# Activate the WiFi interface
|
||||
wlan.active(True)
|
||||
# If not already connected
|
||||
if not wlan.isconnected():
|
||||
# Scan for WiFi networks
|
||||
networks = wlan.scan()
|
||||
# Go trough all scanned WiFi networks
|
||||
for network in networks:
|
||||
# Extract the networks SSID
|
||||
temp_ssid = network[0].decode("utf-8")
|
||||
# Check if the SSID is in the config file
|
||||
if temp_ssid in config["wifis"].keys():
|
||||
ssid = temp_ssid
|
||||
# Start to connect to the pre-configured network
|
||||
wlan.connect(ssid, config["wifis"][ssid])
|
||||
break
|
||||
|
||||
# Clean up
|
||||
import gc
|
||||
gc.collect()
|
19
config.json
19
config.json
@ -1,19 +0,0 @@
|
||||
{
|
||||
"status_led_pin": 5,
|
||||
"battery_coeff": 2.25,
|
||||
"sumo_id": "xxxxxxxx",
|
||||
"firmware_timestamp": "2019.06.09 21:51:00",
|
||||
"firmware_version": "0.7.0",
|
||||
"left_servo_tuning": 33,
|
||||
"right_servo_tuning": 33,
|
||||
"ultrasonic_threshold": 40,
|
||||
"boot_code": "code.py",
|
||||
"left_line_value": 1000,
|
||||
"right_line_value": 1000,
|
||||
"left_line_threshold": 1000,
|
||||
"right_line_threshold": 1000,
|
||||
"sumo_server": "165.227.140.64:443",
|
||||
"wifis": {
|
||||
"RoboKoding": "salakala"
|
||||
}
|
||||
}
|
344
hal.py
344
hal.py
@ -1,344 +0,0 @@
|
||||
import os
|
||||
import ujson
|
||||
from utime import sleep_us, sleep_ms
|
||||
from machine import Pin, PWM, ADC, time_pulse_us
|
||||
|
||||
# LEDs
|
||||
STATUS = 0
|
||||
OPPONENT = 1
|
||||
LEFT_LINE = 2
|
||||
RIGHT_LINE = 3
|
||||
|
||||
# Directions
|
||||
STOP = 0
|
||||
LEFT = 1
|
||||
RIGHT = 2
|
||||
SEARCH = 3
|
||||
FORWARD = 4
|
||||
BACKWARD = 5
|
||||
|
||||
class Sumorobot(object):
|
||||
# Constructor
|
||||
def __init__(self, config = None):
|
||||
# Config file
|
||||
self.config = config
|
||||
|
||||
# Ultrasonic distance sensor
|
||||
self.echo = Pin(14, Pin.IN)
|
||||
self.trigger = Pin(27, Pin.OUT)
|
||||
|
||||
# Servo PWM-s
|
||||
self.pwm_left = PWM(Pin(15), freq=50, duty=0)
|
||||
self.pwm_right = PWM(Pin(4), freq=50, duty=0)
|
||||
|
||||
# LED sensor feedback
|
||||
self.sensor_feedback = True
|
||||
# Bottom status LED
|
||||
self.status_led = Pin(self.config["status_led_pin"], Pin.OUT)
|
||||
# Bottom status LED is in reverse polarity
|
||||
self.status_led.value(1)
|
||||
# Sensor LEDs
|
||||
self.opponent_led = Pin(16, Pin.OUT)
|
||||
self.left_line_led = Pin(17, Pin.OUT)
|
||||
self.right_line_led = Pin(12, Pin.OUT)
|
||||
|
||||
# Scope with sensor data
|
||||
self.sensor_scope = dict()
|
||||
|
||||
# WiFi connection
|
||||
self.is_wifi_connected = False
|
||||
|
||||
# Python and Blockly code
|
||||
self.python_code = ""
|
||||
self.blockly_code = ""
|
||||
self.compiled_python_code = ""
|
||||
|
||||
# Battery gauge
|
||||
self.bat_status = 4.3
|
||||
self.move_counter = 0
|
||||
self.adc_battery = ADC(Pin(32))
|
||||
self.bat_charge = Pin(25, Pin.IN)
|
||||
|
||||
# The pullups for the phototransistors
|
||||
Pin(19, Pin.IN, Pin.PULL_UP)
|
||||
Pin(23, Pin.IN, Pin.PULL_UP)
|
||||
|
||||
# The phototransistors
|
||||
self.last_line = LEFT
|
||||
self.adc_line_left = ADC(Pin(34))
|
||||
self.adc_line_right = ADC(Pin(33))
|
||||
|
||||
# Set reference voltage to 3.3V
|
||||
self.adc_battery.atten(ADC.ATTN_11DB)
|
||||
self.adc_line_left.atten(ADC.ATTN_11DB)
|
||||
self.adc_line_right.atten(ADC.ATTN_11DB)
|
||||
|
||||
# To smooth out ultrasonic sensor value
|
||||
self.opponent_score = 0
|
||||
|
||||
# For terminating sleep
|
||||
self.terminate = False
|
||||
|
||||
# For search mode
|
||||
self.search = False
|
||||
self.search_counter = 0
|
||||
|
||||
# Memorise previous servo speeds
|
||||
self.prev_speed = {LEFT: 0, RIGHT: 0}
|
||||
|
||||
# Function to set LED states
|
||||
def set_led(self, led, state):
|
||||
# Set the given LED state
|
||||
if led == STATUS:
|
||||
# Status LED is reverse polarity
|
||||
self.status_led.value(0 if state else 1)
|
||||
elif led == OPPONENT:
|
||||
self.opponent_led.value(state)
|
||||
elif led == LEFT_LINE:
|
||||
self.left_line_led.value(state)
|
||||
elif led == RIGHT_LINE:
|
||||
self.right_line_led.value(state)
|
||||
|
||||
# Function to shortly bink status LED
|
||||
def toggle_led(self, timer = None):
|
||||
self.status_led.value(0)
|
||||
sleep_ms(10)
|
||||
self.status_led.value(1)
|
||||
|
||||
# Function to get battery voltage
|
||||
def get_battery_voltage(self):
|
||||
bat = round(self.config["battery_coeff"] * (self.adc_battery.read() * 3.3 / 4096), 2)
|
||||
# When the SumoRobot is not moving
|
||||
if self.prev_speed[LEFT] == 0 and self.prev_speed[RIGHT] == 0:
|
||||
if self.move_counter > 0:
|
||||
self.move_counter -= 1
|
||||
if self.bat_status < bat - 0.20 and self.move_counter == 0:
|
||||
#deepsleep()
|
||||
pass
|
||||
self.bat_status = bat
|
||||
else:
|
||||
self.move_counter = 10
|
||||
return bat
|
||||
|
||||
# Function to get distance (cm) from the object in front of the SumoRobot
|
||||
def get_opponent_distance(self):
|
||||
# Send a pulse
|
||||
self.trigger.value(0)
|
||||
sleep_us(5)
|
||||
self.trigger.value(1)
|
||||
sleep_us(10)
|
||||
self.trigger.value(0)
|
||||
# Wait for the pulse and calculate the distance
|
||||
return round((time_pulse_us(self.echo, 1, 30000) / 2) / 29.1, 2)
|
||||
|
||||
# Function to get boolean if there is something in front of the SumoRobot
|
||||
def is_opponent(self):
|
||||
# Get the opponent distance
|
||||
self.opponent_distance = self.get_opponent_distance()
|
||||
# When the opponent is close and the ping actually returned
|
||||
if self.opponent_distance < self.config["ultrasonic_threshold"] and self.opponent_distance > 0:
|
||||
# When not maximum score
|
||||
if self.opponent_score < 5:
|
||||
# Increase the opponent score
|
||||
self.opponent_score += 1
|
||||
# When no opponent was detected
|
||||
else:
|
||||
# When not lowest score
|
||||
if self.opponent_score > 0:
|
||||
# Decrease the opponent score
|
||||
self.opponent_score -= 1
|
||||
|
||||
# When the sensor saw something more than 2 times
|
||||
opponent = True if self.opponent_score > 2 else False
|
||||
|
||||
# Trigger opponent LED
|
||||
self.set_led(OPPONENT, opponent)
|
||||
|
||||
return opponent
|
||||
|
||||
# Function to update the config file
|
||||
def update_config_file(self):
|
||||
# Update the config file
|
||||
with open("config.part", "w") as config_file:
|
||||
config_file.write(ujson.dumps(self.config))
|
||||
os.rename("config.part", "config.json")
|
||||
|
||||
# Function to update line calibration and write it to the config file
|
||||
def calibrate_line_value(self):
|
||||
# Read the line sensor values
|
||||
self.config["left_line_value"] = self.adc_line_left.read()
|
||||
self.config["right_line_value"] = self.adc_line_right.read()
|
||||
# Update the config file
|
||||
self.update_config_file()
|
||||
|
||||
# Function to update line threshold and write it to the config file
|
||||
def set_line_threshold(self, value):
|
||||
# Read the line sensor values
|
||||
self.config["left_line_threshold"] = value
|
||||
self.config["right_line_threshold"] = value
|
||||
# Update the config file
|
||||
self.update_config_file()
|
||||
|
||||
# Function to update ultrasonic sensor threshold and write it to the config file
|
||||
def set_ultrasonic_threshold(self, value):
|
||||
# Read the line sensor values
|
||||
self.config["ultrasonic_threshold"] = value
|
||||
# Update the config file
|
||||
self.update_config_file()
|
||||
|
||||
# Function to get light inensity from the phototransistors
|
||||
def get_line(self, dir):
|
||||
# Check if the direction is valid
|
||||
assert dir in (LEFT, RIGHT)
|
||||
|
||||
# Return the given line sensor value
|
||||
if dir == LEFT:
|
||||
return self.adc_line_left.read()
|
||||
elif dir == RIGHT:
|
||||
return self.adc_line_right.read()
|
||||
|
||||
def is_line(self, dir):
|
||||
# Check if the direction is valid
|
||||
assert dir in (LEFT, RIGHT)
|
||||
|
||||
# Return the given line sensor value
|
||||
if dir == LEFT:
|
||||
line = abs(self.get_line(LEFT) - self.config["left_line_value"]) > self.config["left_line_threshold"]
|
||||
self.set_led(LEFT_LINE, line)
|
||||
last_line = LEFT
|
||||
return line
|
||||
elif dir == RIGHT:
|
||||
line = abs(self.get_line(RIGHT) - self.config["right_line_value"]) > self.config["right_line_threshold"]
|
||||
self.set_led(RIGHT_LINE, line)
|
||||
last_line = RIGHT
|
||||
return line
|
||||
|
||||
def set_servo(self, dir, speed):
|
||||
# Check if the direction is valid
|
||||
assert dir in (LEFT, RIGHT)
|
||||
# Check if the speed is valid
|
||||
assert speed <= 100 and speed >= -100
|
||||
|
||||
# When the speed didn't change
|
||||
if speed == self.prev_speed[dir]:
|
||||
return
|
||||
|
||||
# Record the new speed
|
||||
self.prev_speed[dir] = speed
|
||||
|
||||
# Set the given servo speed
|
||||
if dir == LEFT:
|
||||
if speed == 0:
|
||||
self.pwm_left.duty(0)
|
||||
else:
|
||||
# -100 ... 100 to 33 .. 102
|
||||
self.pwm_left.duty(int(33 + self.config["left_servo_tuning"] + speed * 33 / 100))
|
||||
elif dir == RIGHT:
|
||||
if speed == 0:
|
||||
self.pwm_right.duty(0)
|
||||
else:
|
||||
# -100 ... 100 to 33 .. 102
|
||||
self.pwm_right.duty(int(33 + self.config["right_servo_tuning"] + speed * 33 / 100))
|
||||
|
||||
def move(self, dir):
|
||||
# Check if the direction is valid
|
||||
assert dir in (SEARCH, STOP, RIGHT, LEFT, BACKWARD, FORWARD)
|
||||
|
||||
# Go to the given direction
|
||||
if dir == STOP:
|
||||
self.set_servo(LEFT, 0)
|
||||
self.set_servo(RIGHT, 0)
|
||||
elif dir == LEFT:
|
||||
self.set_servo(LEFT, -100)
|
||||
self.set_servo(RIGHT, -100)
|
||||
elif dir == RIGHT:
|
||||
self.set_servo(LEFT, 100)
|
||||
self.set_servo(RIGHT, 100)
|
||||
elif dir == SEARCH:
|
||||
# Change search mode after X seconds
|
||||
if self.search_counter == 50:
|
||||
self.search = not self.search
|
||||
self.search_counter = 0
|
||||
# When in search mode
|
||||
if self.search:
|
||||
# Go forward
|
||||
self.set_servo(LEFT, 100)
|
||||
self.set_servo(RIGHT, -100)
|
||||
elif last_line == RIGHT:
|
||||
# Go left
|
||||
self.set_servo(LEFT, -100)
|
||||
self.set_servo(RIGHT, -100)
|
||||
else:
|
||||
# Go right
|
||||
self.set_servo(LEFT, 100)
|
||||
self.set_servo(RIGHT, 100)
|
||||
# Increase search counter
|
||||
self.search_counter += 1
|
||||
elif dir == FORWARD:
|
||||
self.set_servo(LEFT, 100)
|
||||
self.set_servo(RIGHT, -100)
|
||||
elif dir == BACKWARD:
|
||||
self.set_servo(LEFT, -100)
|
||||
self.set_servo(RIGHT, 100)
|
||||
|
||||
def update_sensor_feedback(self):
|
||||
if self.sensor_feedback:
|
||||
# Execute to see LED feedback for sensors
|
||||
self.is_opponent()
|
||||
self.is_line(LEFT)
|
||||
self.is_line(RIGHT)
|
||||
|
||||
def update_sensor_scope(self):
|
||||
self.sensor_scope = dict(
|
||||
left_line = self.get_line(LEFT),
|
||||
right_line = self.get_line(RIGHT),
|
||||
opponent = self.get_opponent_distance(),
|
||||
battery_charge = self.bat_charge.value(),
|
||||
battery_voltage = self.get_battery_voltage()
|
||||
)
|
||||
|
||||
def get_python_code(self):
|
||||
return dict(
|
||||
type = "python_code",
|
||||
val = self.python_code
|
||||
)
|
||||
|
||||
def get_blockly_code(self):
|
||||
return dict(
|
||||
type = "blockly_code",
|
||||
val = self.blockly_code
|
||||
)
|
||||
|
||||
def get_firmware_version(self):
|
||||
return dict(
|
||||
type = "firmware_version",
|
||||
val = self.config["firmware_version"]
|
||||
)
|
||||
|
||||
def get_sensor_scope(self):
|
||||
temp = self.sensor_scope
|
||||
temp['type'] = "sensor_scope"
|
||||
return temp
|
||||
|
||||
def get_threshold_scope(self):
|
||||
return dict(
|
||||
type = "threshold_scope",
|
||||
left_line_value = self.config["left_line_value"],
|
||||
right_line_value = self.config["right_line_value"],
|
||||
left_line_threshold = self.config["left_line_threshold"],
|
||||
right_line_threshold = self.config["right_line_threshold"],
|
||||
ultrasonic_threshold = self.config["ultrasonic_threshold"]
|
||||
)
|
||||
|
||||
def sleep(self, delay):
|
||||
# Check for valid delay
|
||||
assert delay > 0
|
||||
|
||||
# Split the delay into 50ms chunks
|
||||
for j in range(0, delay, 50):
|
||||
# Check for forceful termination
|
||||
if self.terminate:
|
||||
# Terminate the delay
|
||||
return
|
||||
else:
|
||||
sleep_ms(50)
|
171
main.py
171
main.py
@ -1,171 +0,0 @@
|
||||
# The code processing thread
|
||||
def step():
|
||||
while True:
|
||||
# Execute to see LED feedback for sensors
|
||||
sumorobot.update_sensor_feedback()
|
||||
# Update sensor scope
|
||||
sumorobot.update_sensor_scope()
|
||||
# Try to execute the Python code
|
||||
try:
|
||||
exec(sumorobot.compiled_python_code)
|
||||
except:
|
||||
pass
|
||||
# When robot was stopped
|
||||
if sumorobot.terminate:
|
||||
# Disable forceful termination of delays in code
|
||||
sumorobot.terminate = False
|
||||
# Stop the robot
|
||||
sumorobot.move(STOP)
|
||||
# Leave time to process WebSocket commands
|
||||
sleep_ms(50)
|
||||
|
||||
# The WebSocket processing thread
|
||||
def ws_handler():
|
||||
global conn, watchdog_counter
|
||||
|
||||
while True:
|
||||
# When WiFi has just been reconnected
|
||||
if wlan.isconnected() and not sumorobot.is_wifi_connected:
|
||||
print("main.py reconnected to Wi-Fi")
|
||||
# Stop blinking status LED
|
||||
timer.deinit()
|
||||
# Turn status LED to steady ON
|
||||
sumorobot.set_led(STATUS, True)
|
||||
sumorobot.is_wifi_connected = True
|
||||
# When WiFi has just been disconnected
|
||||
elif not wlan.isconnected() and sumorobot.is_wifi_connected:
|
||||
print("main.py lost Wi-Fi, reconnecting to Wi-Fi")
|
||||
# Reinitiate the Wi-Fi connection
|
||||
wlan.connect(ssid, config["wifis"][ssid])
|
||||
# Turn OFF status LED
|
||||
sumorobot.set_led(STATUS, False)
|
||||
sumorobot.is_wifi_connected = False
|
||||
# Start bliking status LED
|
||||
timer.init(period=2000, mode=Timer.PERIODIC, callback=sumorobot.toggle_led)
|
||||
elif not wlan.isconnected():
|
||||
# Continue to wait for a WiFi connection
|
||||
continue
|
||||
|
||||
data = None
|
||||
try: # Try to read from the WebSocket
|
||||
data = conn.recv()
|
||||
except: # Socket timeout, no data received
|
||||
# Increment watchdog counter
|
||||
watchdog_counter += 1
|
||||
# When Wi-Fi is connected and X-th exception happened
|
||||
# Try reconnecting to the WebSocket server
|
||||
if wlan.isconnected() and watchdog_counter == 3:
|
||||
print("main.py WebSocket timeout, reconnecting")
|
||||
conn = uwebsockets.connect(uri)
|
||||
watchdog_counter = 0
|
||||
# Continue to try to read data
|
||||
continue
|
||||
|
||||
# When an empty frame was received
|
||||
if not data:
|
||||
# Continue to receive data
|
||||
continue
|
||||
elif b'forward' in data:
|
||||
sumorobot.compiled_python_code = ""
|
||||
sumorobot.move(FORWARD)
|
||||
elif b'backward' in data:
|
||||
sumorobot.compiled_python_code = ""
|
||||
sumorobot.move(BACKWARD)
|
||||
elif b'right' in data:
|
||||
sumorobot.compiled_python_code = ""
|
||||
sumorobot.move(RIGHT)
|
||||
elif b'left' in data:
|
||||
sumorobot.compiled_python_code = ""
|
||||
sumorobot.move(LEFT)
|
||||
elif b'stop' in data:
|
||||
sumorobot.compiled_python_code = ""
|
||||
sumorobot.move(STOP)
|
||||
# for terminating delays in code
|
||||
sumorobot.terminate = True
|
||||
elif b'get_threshold_scope' in data:
|
||||
conn.send(ujson.dumps(sumorobot.get_threshold_scope()))
|
||||
elif b'get_sensor_scope' in data:
|
||||
conn.send(ujson.dumps(sumorobot.get_sensor_scope()))
|
||||
elif b'get_python_code' in data:
|
||||
#print("main.py sending python code=", sumorobot.get_python_code())
|
||||
conn.send(ujson.dumps(sumorobot.get_python_code()))
|
||||
elif b'get_blockly_code' in data:
|
||||
#print("main.py sending blockly code=", sumorobot.get_blockly_code())
|
||||
conn.send(ujson.dumps(sumorobot.get_blockly_code()))
|
||||
elif b'get_firmware_version' in data:
|
||||
#print("main.py get_firmware_version")
|
||||
conn.send(ujson.dumps(sumorobot.get_firmware_version()))
|
||||
elif b'toggle_sensor_feedback' in data:
|
||||
data = ujson.loads(data)
|
||||
sumorobot.sensor_feedback = not sumorobot.sensor_feedback
|
||||
elif b'set_blockly_code' in data:
|
||||
data = ujson.loads(data)
|
||||
#print("main.py Blockly code=", data['val'])
|
||||
sumorobot.blockly_code = data['val']
|
||||
elif b'set_python_code' in data:
|
||||
data = ujson.loads(data)
|
||||
sumorobot.python_code = data['val']
|
||||
data['val'] = data['val'].replace(";;", "\n")
|
||||
#print("main.py python code=", data['val'])
|
||||
sumorobot.compiled_python_code = compile(data['val'], "snippet", "exec")
|
||||
elif b'calibrate_line_value' in data:
|
||||
sumorobot.calibrate_line_value()
|
||||
#print("main.py: calibrate_line_value")
|
||||
elif b'set_line_threshold' in data:
|
||||
data = ujson.loads(data)
|
||||
sumorobot.set_line_threshold(int(data['val']))
|
||||
#print("main.py: set_line_threshold")
|
||||
elif b'set_ultrasonic_threshold' in data:
|
||||
data = ujson.loads(data)
|
||||
sumorobot.set_ultrasonic_threshold(int(data['val']))
|
||||
#print("main.py: set_ultrasonic_threshold")
|
||||
elif b'Gone' in data:
|
||||
print("main.py: server said 410 Gone, attempting to reconnect...")
|
||||
#conn = uwebsockets.connect(url)
|
||||
else:
|
||||
print("main.py: unknown cmd=", data)
|
||||
|
||||
# When user code (copy.py) exists
|
||||
if 'code.py' in os.listdir():
|
||||
print("main.py: loading user code")
|
||||
# Try to load the user code
|
||||
try:
|
||||
with open("code.py", "r") as code:
|
||||
temp = code.read()
|
||||
sumorobot.python_code = temp
|
||||
sumorobot.compiled_python_code = compile(temp, "snippet", "exec")
|
||||
except:
|
||||
print("main.py: error loading code.py file")
|
||||
|
||||
# Start the code processing thread
|
||||
_thread.start_new_thread(step, ())
|
||||
|
||||
# Wifi watchdog counter
|
||||
watchdog_counter = 0
|
||||
# Wait for WiFi to get connected
|
||||
while not wlan.isconnected():
|
||||
sleep_ms(100)
|
||||
watchdog_counter += 1
|
||||
# When Wi-Fi didn't connect in X seconds
|
||||
if watchdog_counter == 30:
|
||||
print("main.py reconnecting to Wi-Fi")
|
||||
# Reinitiate the Wi-Fi connection
|
||||
wlan.connect(ssid, config["wifis"][ssid])
|
||||
|
||||
# Restart watchdog counter
|
||||
watchdog_counter = 0
|
||||
|
||||
# Connect to the websocket
|
||||
uri = "ws://%s/p2p/sumo-%s/browser/" % (config['sumo_server'], config['sumo_id'])
|
||||
conn = uwebsockets.connect(uri)
|
||||
|
||||
# Stop bootup blinking
|
||||
timer.deinit()
|
||||
|
||||
# WiFi is connected
|
||||
sumorobot.is_wifi_connected = True
|
||||
# Indicate that the WebSocket is connected
|
||||
sumorobot.set_led(STATUS, True)
|
||||
|
||||
# Start the Websocket processing thread
|
||||
_thread.start_new_thread(ws_handler, ())
|
17
platformio.ini
Normal file
17
platformio.ini
Normal file
@ -0,0 +1,17 @@
|
||||
;PlatformIO Project Configuration File
|
||||
;
|
||||
; Build options: build flags, source filter
|
||||
; Upload options: custom upload port, speed and extra flags
|
||||
; Library options: dependencies, extra library storages
|
||||
; Advanced options: extra scripting
|
||||
;
|
||||
; Please visit documentation for the other options and examples
|
||||
; https://docs.platformio.org/page/projectconf.html
|
||||
|
||||
[env:lolin32]
|
||||
platform = espressif32
|
||||
board = lolin32
|
||||
framework = arduino
|
||||
upload_speed = 460800
|
||||
monitor_speed = 115200
|
||||
lib_deps = NewPing
|
396
src/main.cpp
Normal file
396
src/main.cpp
Normal file
@ -0,0 +1,396 @@
|
||||
/*
|
||||
This the code that runs on the SumoRobots
|
||||
*/
|
||||
// Include BLE libraries
|
||||
#include <BLEDevice.h>
|
||||
#include <BLEServer.h>
|
||||
#include <BLEUtils.h>
|
||||
#include <BLE2902.h>
|
||||
#include <BLE2904.h>
|
||||
// Include other libraries
|
||||
#include <string.h>
|
||||
#include <Ticker.h>
|
||||
#include <NewPing.h>
|
||||
#include <Arduino.h>
|
||||
#include <Preferences.h>
|
||||
|
||||
#define DEBUG true
|
||||
|
||||
#define VERSION "0.8.0"
|
||||
#define VERSION_TIMESTAMP "2019.08.13 08:00"
|
||||
|
||||
// See the following for generating UUIDs:
|
||||
// https://www.uuidgenerator.net/
|
||||
#define NUS_SERVICE_UUID "6E400001-B5A3-F393-E0A9-E50E24DCCA9E" // NUS service UUID
|
||||
#define NUS_CHARACTERISTIC_RX_UUID "6E400002-B5A3-F393-E0A9-E50E24DCCA9E"
|
||||
#define NUS_CHARACTERISTIC_TX_UUID "6E400003-B5A3-F393-E0A9-E50E24DCCA9E"
|
||||
|
||||
// Cleate BLE variables
|
||||
BLEServer * bleServer = NULL;
|
||||
bool deviceConnected = false;
|
||||
bool oldDeviceConnected = false;
|
||||
BLECharacteristic * nusTxCharacteristic;
|
||||
BLECharacteristic * batteryLevelCharacteristic;
|
||||
|
||||
// Create preferences persistence
|
||||
Preferences preferences;
|
||||
|
||||
// Create timers
|
||||
Ticker sonarTimer;
|
||||
Ticker batteryTimer;
|
||||
Ticker connectionLedTimer;
|
||||
|
||||
// Battery stuff
|
||||
float batteryVoltage;
|
||||
bool robotMoving = false;
|
||||
uint8_t batteryLevel = 0;
|
||||
uint8_t tempBatteryLevel = 0;
|
||||
|
||||
// Sonar stuff
|
||||
uint8_t sonarValue;
|
||||
NewPing sonar(27, 14, 200);
|
||||
uint8_t sonarThreshold = 40;
|
||||
|
||||
// Line stuff
|
||||
uint16_t leftLineValue;
|
||||
uint16_t rightLineValue;
|
||||
uint16_t leftLineValueField = 0;
|
||||
uint16_t rightLineValueField = 0;
|
||||
uint16_t leftLineThreshold = 1000;
|
||||
uint16_t rightLineThreshold = 1000;
|
||||
|
||||
// Other sensor stuff
|
||||
uint8_t sensorValues[6];
|
||||
bool ledFeedbackEnabled = true;
|
||||
|
||||
// Move command names
|
||||
std::string cmdStop("stop");
|
||||
std::string cmdLeft("left");
|
||||
std::string cmdRight("right");
|
||||
std::string cmdForward("forward");
|
||||
std::string cmdBackward("backward");
|
||||
// Other command names
|
||||
std::string cmdLed("led");
|
||||
std::string cmdLine("line");
|
||||
std::string cmdName("name");
|
||||
std::string cmdSonar("sonar");
|
||||
std::string cmdServo("servo");
|
||||
std::string cmdLedFeedback("ledf");
|
||||
|
||||
void setLed(char led, char value) {
|
||||
// Convert the value to a HIGH or LOW
|
||||
bool state = value == '1' ? HIGH : LOW;
|
||||
if (led == 'c') {
|
||||
// Connection status LED is opposite value
|
||||
digitalWrite(5, !state);
|
||||
}
|
||||
else if (led == 's') {
|
||||
digitalWrite(16, state);
|
||||
}
|
||||
else if (led == 'r') {
|
||||
digitalWrite(12, state);
|
||||
}
|
||||
else if (led == 'l') {
|
||||
digitalWrite(17, state);
|
||||
}
|
||||
}
|
||||
|
||||
void setServo(char servo, int8_t speed) {
|
||||
Serial.println(speed);
|
||||
if (servo == 'l') {
|
||||
ledcWrite(1, map(speed, -100, 100, 1, 100));
|
||||
}
|
||||
else if (servo == 'r') {
|
||||
ledcWrite(2, map(speed, -100, 100, 1, 30));
|
||||
}
|
||||
}
|
||||
|
||||
void updateSensorFeedback() {
|
||||
if (sonarValue <= sonarThreshold) {
|
||||
digitalWrite(16, HIGH);
|
||||
}
|
||||
else {
|
||||
digitalWrite(16, LOW);
|
||||
}
|
||||
if (abs(leftLineValue - leftLineValueField) > leftLineThreshold) {
|
||||
digitalWrite(17, HIGH);
|
||||
}
|
||||
else {
|
||||
digitalWrite(17, LOW);
|
||||
}
|
||||
if (abs(rightLineValue - rightLineValueField) > rightLineThreshold) {
|
||||
digitalWrite(12, HIGH);
|
||||
}
|
||||
else {
|
||||
digitalWrite(12, LOW);
|
||||
}
|
||||
}
|
||||
|
||||
void updateSonarValue() {
|
||||
// Update the sensor values
|
||||
sonarValue = sonar.ping_cm();
|
||||
// When we didn't receive a ping back
|
||||
// set to max distance
|
||||
if (sonarValue == 0) sonarValue = 255;
|
||||
leftLineValue = analogRead(34);
|
||||
rightLineValue = analogRead(33);
|
||||
if (ledFeedbackEnabled) updateSensorFeedback();
|
||||
sensorValues[0] = sonarValue;
|
||||
sensorValues[1] = leftLineValue >> 8;
|
||||
sensorValues[2] = leftLineValue;
|
||||
sensorValues[3] = rightLineValue >> 8;
|
||||
sensorValues[4] = rightLineValue;
|
||||
sensorValues[5] = digitalRead(25);
|
||||
// When BLE is connected
|
||||
if (deviceConnected) {
|
||||
// Notify the new sensor values
|
||||
nusTxCharacteristic->setValue(sensorValues, 6);
|
||||
nusTxCharacteristic->notify();
|
||||
}
|
||||
}
|
||||
|
||||
void updateBatteryLevel() {
|
||||
// Don't update battery level when robot is moving
|
||||
// the servo motors lower the battery voltage
|
||||
// TODO: wait still a little more after moving
|
||||
// for the voltage to settle
|
||||
if (robotMoving) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Calculate the battery voltage
|
||||
batteryVoltage = 2.12 * (analogRead(32) * 3.3 / 4096);
|
||||
// Calculate battery percentage
|
||||
tempBatteryLevel = 0.0 + ((100.0 - 0.0) / (4.2 - 3.2)) * (batteryVoltage - 3.2);
|
||||
// When battery level changed more than 3%
|
||||
if (abs(batteryLevel - tempBatteryLevel) > 3) {
|
||||
// Update battery level
|
||||
batteryLevel = tempBatteryLevel;
|
||||
}
|
||||
// Notify the new battery level
|
||||
batteryLevelCharacteristic->setValue(&batteryLevel, 1);
|
||||
batteryLevelCharacteristic->notify();
|
||||
#if DEBUG
|
||||
Serial.print(batteryVoltage);
|
||||
Serial.print(" : ");
|
||||
Serial.println(batteryLevel);
|
||||
#endif
|
||||
}
|
||||
|
||||
void blinkConnectionLed() {
|
||||
digitalWrite(5, LOW);
|
||||
delay(20);
|
||||
digitalWrite(5, HIGH);
|
||||
}
|
||||
|
||||
// BLE connect and disconnect callbacks
|
||||
class MyServerCallbacks: public BLEServerCallbacks {
|
||||
void onConnect(BLEServer * pServer) {
|
||||
deviceConnected = true;
|
||||
};
|
||||
|
||||
void onDisconnect(BLEServer * pServer) {
|
||||
deviceConnected = false;
|
||||
}
|
||||
};
|
||||
|
||||
// BLE NUS received callback
|
||||
class MyCallbacks: public BLECharacteristicCallbacks {
|
||||
void onWrite(BLECharacteristic * nusRxCharacteristic) {
|
||||
// Get the received command over BLE
|
||||
std::string cmd = nusRxCharacteristic->getValue();
|
||||
#if DEBUG
|
||||
Serial.println(cmd.c_str());
|
||||
#endif
|
||||
if (cmd.length() > 0) {
|
||||
//int speed = atoi(rxValue.c_str());
|
||||
//Serial.println(speed);
|
||||
//ledcWrite(1, speed); // left 1 ... 100
|
||||
//ledcWrite(2, speed); // right 1 ... 30
|
||||
// Specify command
|
||||
if (cmd == cmdForward) {
|
||||
robotMoving = true;
|
||||
ledcWrite(1, 100);
|
||||
ledcWrite(2, 1);
|
||||
}
|
||||
else if (cmd == cmdBackward) {
|
||||
robotMoving = true;
|
||||
ledcWrite(1, 1);
|
||||
ledcWrite(2, 30);
|
||||
}
|
||||
else if (cmd == cmdLeft) {
|
||||
robotMoving = true;
|
||||
ledcWrite(1, 1);
|
||||
ledcWrite(2, 1);
|
||||
}
|
||||
else if (cmd == cmdRight) {
|
||||
robotMoving = true;
|
||||
ledcWrite(1, 100);
|
||||
ledcWrite(2, 30);
|
||||
}
|
||||
else if (cmd == cmdStop) {
|
||||
robotMoving = false;
|
||||
ledcWrite(1, 0);
|
||||
ledcWrite(2, 0);
|
||||
}
|
||||
else if (cmd == cmdLedFeedback) {
|
||||
ledFeedbackEnabled = !ledFeedbackEnabled;
|
||||
}
|
||||
else if (cmd.find(cmdLed) != std::string::npos) {
|
||||
setLed(cmd.at(3), cmd.at(4));
|
||||
}
|
||||
else if (cmd.find(cmdLine) != std::string::npos) {
|
||||
// Get the threshold value
|
||||
leftLineThreshold = atoi(cmd.substr(4, cmd.length() - 4).c_str());
|
||||
rightLineThreshold = leftLineThreshold;
|
||||
// Remember value on the field (white or black)
|
||||
leftLineValueField = analogRead(34);
|
||||
rightLineValueField = analogRead(33);
|
||||
// Save the threshold value in the persistence
|
||||
preferences.begin("sumorobot", false);
|
||||
preferences.putUInt("line_threshold", leftLineThreshold);
|
||||
preferences.end();
|
||||
}
|
||||
else if (cmd.find(cmdSonar) != std::string::npos) {
|
||||
sonarThreshold = atoi(cmd.substr(5, cmd.length() - 5).c_str());
|
||||
// Save the threshold value in the persistence
|
||||
preferences.begin("sumorobot", false);
|
||||
preferences.putUInt("sonar_threshold", sonarThreshold);
|
||||
preferences.end();
|
||||
}
|
||||
else if (cmd.find(cmdServo) != std::string::npos) {
|
||||
setServo(cmd.at(5), atoi(cmd.substr(6, cmd.length() - 6).c_str()));
|
||||
}
|
||||
else if (cmd.find(cmdName) != std::string::npos) {
|
||||
preferences.begin("sumorobot", false);
|
||||
preferences.putString("name", cmd.substr(4, cmd.length() - 4).c_str());
|
||||
preferences.end();
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
void setup() {
|
||||
#if DEBUG
|
||||
Serial.begin(115200);
|
||||
#endif
|
||||
|
||||
// Start preferences persistence
|
||||
preferences.begin("sumorobot", false);
|
||||
|
||||
// Create the BLE device
|
||||
Serial.println(preferences.getString("name", "SumoRobot").c_str());
|
||||
BLEDevice::init(preferences.getString("name", "SumoRobot").c_str());
|
||||
|
||||
preferences.end();
|
||||
|
||||
// Create the BLE server
|
||||
bleServer = BLEDevice::createServer();
|
||||
bleServer->setCallbacks(new MyServerCallbacks());
|
||||
|
||||
// Create device info service and characteristic
|
||||
BLEService * deviceInfoService = bleServer->createService(BLEUUID((uint16_t) 0x180a));
|
||||
BLECharacteristic * modelCharacteristic = deviceInfoService->createCharacteristic(
|
||||
(uint16_t) 0x2A24, BLECharacteristic::PROPERTY_READ);
|
||||
BLECharacteristic * firmwareCharacteristic = deviceInfoService->createCharacteristic(
|
||||
(uint16_t) 0x2A26, BLECharacteristic::PROPERTY_READ);
|
||||
BLECharacteristic * manufacturerCharacteristic = deviceInfoService->createCharacteristic(
|
||||
(uint16_t) 0x2a29, BLECharacteristic::PROPERTY_READ);
|
||||
manufacturerCharacteristic->setValue("RoboKoding LTD");
|
||||
modelCharacteristic->setValue("SumoRobot");
|
||||
firmwareCharacteristic->setValue(VERSION);
|
||||
|
||||
// Create battery service
|
||||
BLEService * batteryService = bleServer->createService(BLEUUID((uint16_t) 0x180f));
|
||||
// Mandatory battery level characteristic with notification and presence descriptor
|
||||
BLE2904* batteryLevelDescriptor = new BLE2904();
|
||||
batteryLevelDescriptor->setFormat(BLE2904::FORMAT_UINT8);
|
||||
batteryLevelDescriptor->setNamespace(1);
|
||||
batteryLevelDescriptor->setUnit(0x27ad);
|
||||
// Create battery level characteristics
|
||||
batteryLevelCharacteristic = batteryService->createCharacteristic(
|
||||
(uint16_t) 0x2a19, BLECharacteristic::PROPERTY_READ | BLECharacteristic::PROPERTY_NOTIFY);
|
||||
batteryLevelCharacteristic->addDescriptor(batteryLevelDescriptor);
|
||||
batteryLevelCharacteristic->addDescriptor(new BLE2902());
|
||||
|
||||
// Create the BLE NUS service
|
||||
BLEService * nusService = bleServer->createService(NUS_SERVICE_UUID);
|
||||
|
||||
// Create a BLE NUS transmit characteristic
|
||||
nusTxCharacteristic = nusService->createCharacteristic(
|
||||
NUS_CHARACTERISTIC_TX_UUID, BLECharacteristic::PROPERTY_NOTIFY);
|
||||
nusTxCharacteristic->addDescriptor(new BLE2902());
|
||||
|
||||
// Create a BLE NUS receive characteristics
|
||||
BLECharacteristic * nusRxCharacteristic = nusService->createCharacteristic(
|
||||
NUS_CHARACTERISTIC_RX_UUID, BLECharacteristic::PROPERTY_WRITE);
|
||||
nusRxCharacteristic->setCallbacks(new MyCallbacks());
|
||||
|
||||
// Start the services
|
||||
deviceInfoService->start();
|
||||
batteryService->start();
|
||||
nusService->start();
|
||||
|
||||
// Start advertising
|
||||
bleServer->getAdvertising()->start();
|
||||
|
||||
#if DEBUG
|
||||
Serial.println("Waiting a client connection to notify...");
|
||||
#endif
|
||||
|
||||
// Setup BLE connection status LED
|
||||
pinMode(5, OUTPUT);
|
||||
connectionLedTimer.attach_ms(2000, blinkConnectionLed);
|
||||
|
||||
// Setup the left servo PWM
|
||||
ledcSetup(1, 50, 10);
|
||||
ledcAttachPin(15, 1);
|
||||
|
||||
// Setup the right servo PWM
|
||||
ledcSetup(2, 50, 8);
|
||||
ledcAttachPin(4, 2);
|
||||
|
||||
// Phototransistor pull-ups
|
||||
pinMode(19, INPUT_PULLUP);
|
||||
pinMode(23, INPUT_PULLUP);
|
||||
|
||||
// Setup battery charge detection pin
|
||||
pinMode(25, INPUT);
|
||||
|
||||
// Setup sensor feedback LED pins
|
||||
pinMode(16, OUTPUT);
|
||||
pinMode(17, OUTPUT);
|
||||
pinMode(12, OUTPUT);
|
||||
|
||||
// Setup ADC for reading phototransistors and battery
|
||||
analogSetAttenuation(ADC_11db);
|
||||
adcAttachPin(32);
|
||||
adcAttachPin(33);
|
||||
adcAttachPin(34);
|
||||
|
||||
// Setup sonar timer to update it's value
|
||||
sonarTimer.attach_ms(50, updateSonarValue);
|
||||
|
||||
// Setup battery level timer to update it's value
|
||||
batteryTimer.attach(5, updateBatteryLevel);
|
||||
updateBatteryLevel();
|
||||
}
|
||||
|
||||
void loop() {
|
||||
// When BLE got disconnected
|
||||
if (!deviceConnected && oldDeviceConnected) {
|
||||
delay(500); // Give the bluetooth stack the chance to get things ready
|
||||
bleServer->startAdvertising(); // Restart advertising
|
||||
#if DEBUG
|
||||
Serial.println("start advertising");
|
||||
#endif
|
||||
oldDeviceConnected = deviceConnected;
|
||||
connectionLedTimer.attach_ms(2000, blinkConnectionLed);
|
||||
}
|
||||
// When BLE got connected
|
||||
if (deviceConnected && !oldDeviceConnected) {
|
||||
oldDeviceConnected = deviceConnected;
|
||||
connectionLedTimer.detach();
|
||||
digitalWrite(5, LOW);
|
||||
}
|
||||
}
|
285
uwebsockets.py
285
uwebsockets.py
@ -1,285 +0,0 @@
|
||||
"""
|
||||
Websockets client for micropython
|
||||
|
||||
Based very heavily off
|
||||
https://github.com/aaugustin/websockets/blob/master/websockets/client.py
|
||||
"""
|
||||
|
||||
#import libraries
|
||||
import ussl
|
||||
import ure as re
|
||||
import urandom as random
|
||||
import ustruct as struct
|
||||
import usocket as socket
|
||||
import ubinascii as binascii
|
||||
from ucollections import namedtuple
|
||||
|
||||
# Opcodes
|
||||
OP_CONT = const(0x0)
|
||||
OP_TEXT = const(0x1)
|
||||
OP_BYTES = const(0x2)
|
||||
OP_CLOSE = const(0x8)
|
||||
OP_PING = const(0x9)
|
||||
OP_PONG = const(0xa)
|
||||
|
||||
# Close codes
|
||||
CLOSE_OK = const(1000)
|
||||
CLOSE_GOING_AWAY = const(1001)
|
||||
CLOSE_PROTOCOL_ERROR = const(1002)
|
||||
CLOSE_DATA_NOT_SUPPORTED = const(1003)
|
||||
CLOSE_BAD_DATA = const(1007)
|
||||
CLOSE_POLICY_VIOLATION = const(1008)
|
||||
CLOSE_TOO_BIG = const(1009)
|
||||
CLOSE_MISSING_EXTN = const(1010)
|
||||
CLOSE_BAD_CONDITION = const(1011)
|
||||
|
||||
URL_RE = re.compile(r'(wss|ws)://([A-Za-z0-9-\.]+)(?:\:([0-9]+))?(/.+)?')
|
||||
URI = namedtuple('URI', ('protocol', 'hostname', 'port', 'path'))
|
||||
|
||||
class NoDataException(Exception):
|
||||
pass
|
||||
|
||||
def urlparse(uri):
|
||||
"""Parse ws:// URLs"""
|
||||
match = URL_RE.match(uri)
|
||||
if match:
|
||||
protocol = match.group(1)
|
||||
host = match.group(2)
|
||||
port = match.group(3)
|
||||
path = match.group(4)
|
||||
|
||||
if protocol == 'wss':
|
||||
if port is None:
|
||||
port = 443
|
||||
elif protocol == 'ws':
|
||||
if port is None:
|
||||
port = 80
|
||||
else:
|
||||
raise ValueError('Scheme {} is invalid'.format(protocol))
|
||||
|
||||
return URI(protocol, host, int(port), path)
|
||||
|
||||
|
||||
class Websocket:
|
||||
"""
|
||||
Basis of the Websocket protocol.
|
||||
|
||||
This can probably be replaced with the C-based websocket module, but
|
||||
this one currently supports more options.
|
||||
"""
|
||||
is_client = False
|
||||
|
||||
def __init__(self, sock):
|
||||
self.sock = sock
|
||||
self.open = True
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
self.close()
|
||||
|
||||
def settimeout(self, timeout):
|
||||
self.sock.settimeout(timeout)
|
||||
|
||||
def read_frame(self, max_size=None):
|
||||
"""
|
||||
Read a frame from the socket.
|
||||
See https://tools.ietf.org/html/rfc6455#section-5.2 for the details.
|
||||
"""
|
||||
|
||||
# Frame header
|
||||
two_bytes = self.sock.read(2)
|
||||
|
||||
if not two_bytes:
|
||||
raise NoDataException
|
||||
|
||||
byte1, byte2 = struct.unpack('!BB', two_bytes)
|
||||
|
||||
# Byte 1: FIN(1) _(1) _(1) _(1) OPCODE(4)
|
||||
fin = bool(byte1 & 0x80)
|
||||
opcode = byte1 & 0x0f
|
||||
|
||||
# Byte 2: MASK(1) LENGTH(7)
|
||||
mask = bool(byte2 & (1 << 7))
|
||||
length = byte2 & 0x7f
|
||||
|
||||
if length == 126: # Magic number, length header is 2 bytes
|
||||
length, = struct.unpack('!H', self.sock.read(2))
|
||||
elif length == 127: # Magic number, length header is 8 bytes
|
||||
length, = struct.unpack('!Q', self.sock.read(8))
|
||||
|
||||
if mask: # Mask is 4 bytes
|
||||
mask_bits = self.sock.read(4)
|
||||
|
||||
try:
|
||||
data = self.sock.read(length)
|
||||
except MemoryError:
|
||||
# We can't receive this many bytes, close the socket
|
||||
self.close(code=CLOSE_TOO_BIG)
|
||||
return True, OP_CLOSE, None
|
||||
|
||||
if mask:
|
||||
data = bytes(b ^ mask_bits[i % 4]
|
||||
for i, b in enumerate(data))
|
||||
|
||||
return fin, opcode, data
|
||||
|
||||
def write_frame(self, opcode, data=b''):
|
||||
"""
|
||||
Write a frame to the socket.
|
||||
See https://tools.ietf.org/html/rfc6455#section-5.2 for the details.
|
||||
"""
|
||||
fin = True
|
||||
mask = self.is_client # messages sent by client are masked
|
||||
|
||||
length = len(data)
|
||||
|
||||
# Frame header
|
||||
# Byte 1: FIN(1) _(1) _(1) _(1) OPCODE(4)
|
||||
byte1 = 0x80 if fin else 0
|
||||
byte1 |= opcode
|
||||
|
||||
# Byte 2: MASK(1) LENGTH(7)
|
||||
byte2 = 0x80 if mask else 0
|
||||
|
||||
if length < 126: # 126 is magic value to use 2-byte length header
|
||||
byte2 |= length
|
||||
self.sock.write(struct.pack('!BB', byte1, byte2))
|
||||
|
||||
elif length < (1 << 16): # Length fits in 2-bytes
|
||||
byte2 |= 126 # Magic code
|
||||
self.sock.write(struct.pack('!BBH', byte1, byte2, length))
|
||||
|
||||
elif length < (1 << 64):
|
||||
byte2 |= 127 # Magic code
|
||||
self.sock.write(struct.pack('!BBQ', byte1, byte2, length))
|
||||
|
||||
else:
|
||||
raise ValueError()
|
||||
|
||||
if mask: # Mask is 4 bytes
|
||||
mask_bits = struct.pack('!I', random.getrandbits(32))
|
||||
self.sock.write(mask_bits)
|
||||
|
||||
data = bytes(b ^ mask_bits[i % 4]
|
||||
for i, b in enumerate(data))
|
||||
|
||||
self.sock.write(data)
|
||||
|
||||
def recv(self):
|
||||
"""
|
||||
Receive data from the websocket.
|
||||
|
||||
This is slightly different from 'websockets' in that it doesn't
|
||||
fire off a routine to process frames and put the data in a queue.
|
||||
If you don't call recv() sufficiently often you won't process control
|
||||
frames.
|
||||
"""
|
||||
assert self.open
|
||||
|
||||
while self.open:
|
||||
try:
|
||||
fin, opcode, data = self.read_frame()
|
||||
except NoDataException:
|
||||
return ''
|
||||
except ValueError:
|
||||
self._close()
|
||||
return
|
||||
|
||||
if not fin:
|
||||
raise NotImplementedError()
|
||||
|
||||
if opcode == OP_TEXT:
|
||||
return data
|
||||
elif opcode == OP_BYTES:
|
||||
return data
|
||||
elif opcode == OP_CLOSE:
|
||||
self._close()
|
||||
return
|
||||
elif opcode == OP_PONG:
|
||||
# Ignore this frame, keep waiting for a data frame
|
||||
continue
|
||||
elif opcode == OP_PING:
|
||||
# We need to send a pong frame
|
||||
self.write_frame(OP_PONG, data)
|
||||
# And then wait to receive
|
||||
continue
|
||||
elif opcode == OP_CONT:
|
||||
# This is a continuation of a previous frame
|
||||
raise NotImplementedError(opcode)
|
||||
else:
|
||||
raise ValueError(opcode)
|
||||
|
||||
def send(self, buf):
|
||||
"""Send data to the websocket."""
|
||||
|
||||
assert self.open
|
||||
|
||||
if isinstance(buf, str):
|
||||
opcode = OP_TEXT
|
||||
buf = buf.encode('utf-8')
|
||||
elif isinstance(buf, bytes):
|
||||
opcode = OP_BYTES
|
||||
else:
|
||||
raise TypeError()
|
||||
|
||||
self.write_frame(opcode, buf)
|
||||
|
||||
def close(self, code=CLOSE_OK, reason=''):
|
||||
"""Close the websocket."""
|
||||
if not self.open:
|
||||
return
|
||||
|
||||
buf = struct.pack('!H', code) + reason.encode('utf-8')
|
||||
|
||||
self.write_frame(OP_CLOSE, buf)
|
||||
self._close()
|
||||
|
||||
def _close(self):
|
||||
self.open = False
|
||||
self.sock.close()
|
||||
|
||||
class WebsocketClient(Websocket):
|
||||
is_client = True
|
||||
|
||||
def connect(uri):
|
||||
"""
|
||||
Connect a websocket.
|
||||
"""
|
||||
|
||||
uri = urlparse(uri)
|
||||
assert uri
|
||||
|
||||
sock = socket.socket()
|
||||
addr = socket.getaddrinfo(uri.hostname, uri.port)
|
||||
sock.connect(addr[0][4])
|
||||
if uri.protocol == 'wss':
|
||||
sock = ussl.wrap_socket(sock)
|
||||
|
||||
def send_header(header, *args):
|
||||
sock.write(header % args + '\r\n')
|
||||
|
||||
# Sec-WebSocket-Key is 16 bytes of random base64 encoded
|
||||
key = binascii.b2a_base64(bytes(random.getrandbits(8)
|
||||
for _ in range(16)))[:-1]
|
||||
|
||||
send_header(b'GET %s HTTP/1.1', uri.path or '/')
|
||||
send_header(b'Host: %s:%s', uri.hostname, uri.port)
|
||||
send_header(b'Connection: Upgrade')
|
||||
send_header(b'Upgrade: websocket')
|
||||
send_header(b'Sec-WebSocket-Key: %s', key)
|
||||
send_header(b'Sec-WebSocket-Version: 13')
|
||||
send_header(b'Origin: http://localhost')
|
||||
send_header(b'')
|
||||
|
||||
header = sock.readline()[:-2]
|
||||
assert header.startswith(b'HTTP/1.1 101 '), header
|
||||
|
||||
# We don't (currently) need these headers
|
||||
# FIXME: should we check the return key?
|
||||
while header:
|
||||
print("uwebsockets.py header:", header)
|
||||
header = sock.readline()[:-2]
|
||||
|
||||
return WebsocketClient(sock)
|
Loading…
Reference in New Issue
Block a user