#!/usr/bin/python3
# This program connects to an Espressif micro-controller (ESP32 or ESP8266) and finds out with which controller
# it is communicating. Then it selects the corresponding Micro Python firmware and downloads it from github
# It uses my github repository for uPyCraft, where the firmware is stored in the directory microPythonFirmware
# After download it checks the file's sha256 hash code to make sure that the file is not corrupted.
# Then the micro-controller flash is erased and the new version of the firmware uploaded and flashed.
# Finally the micro-controller is reset
# The esptool.py program is used for communication with the micro-controller
# Uli Raich 27.2.2019

import esptool
import subprocess,shlex
import hashlib
import sys
import os
import serial
import shutil
import urllib.request
from time import sleep

BUF_SIZE=BUF_SIZE = 65536  # lets read stuff in 64kb chunks!

# read from pipe in chunks of 20 chars

CHUNK_SIZE=20
def each_chunk(stream, separator):
  buffer = ''
  while True:  # until EOF
    chunk = stream.read(CHUNK_SIZE)  # I propose 4096 or so
    if not chunk:  # EOF?
      yield buffer
      break
    buffer += chunk.decode("utf-8")
    while True:  # until no separator is found
      try:
        part, buffer = buffer.split(separator, 1)
      except ValueError:
        break
      else:
        yield part
    
# check if the serial port is connected
try:
    ser = serial.Serial(port='/dev/ttyUSB0')
except:
    print("Could not open serial port /dev/ttyUSB0. Is the micro controller connected?")
    sys.exit()
    
firmwareESP32='/opt/ucc/micros/esp32/micropython/ports/esp32/build/firmware.bin'
firmwareESP8266='/opt/ucc/micros/esp8266/micropython/ports/esp8266/build/firmware-combined.bin'
print("You are using esptool version %s"%esptool.__version__)
progFW=esptool.ESPLoader().detect_chip()

print("Chip ID: %s"%progFW.CHIP_NAME.lower())
if progFW.CHIP_NAME.lower() == "esp8266":
    print("Chip ID: 0x%x"%progFW.chip_id())
elif progFW.CHIP_NAME.lower() == "esp32":
    mac=progFW.read_mac()
    print("%x"%mac[3],end='')
    print("%x"%mac[4],end='')
    print("%x"%mac[5])
    
# get at the flash chip id
proc=subprocess.Popen(['esptool', 'flash_id'],stdout=subprocess.PIPE)

for line in proc.stdout:
    lineString = line.strip().decode("utf-8")
    if 'Chip' in lineString:
        print(lineString)
    if 'Manu' in lineString:
        print("Flash Chip %s"%lineString)
    if 'MAC' in lineString:
        print(lineString)
    if 'flash size' in lineString:
        print(lineString)

print("---- Downloading the micropython firmware ----")

if progFW.CHIP_NAME.lower() == "esp32":
    url = 'https://raw.githubusercontent.com/uraich/uPyCraft-Qt5/master/microPythonFirmware/microPython-ESP32.bin'
    sha256Hash="e22c5329a1bd0d7a3e61b3236a8509a54be8d337629fb12fa3fe028711bd7554"
elif progFW.CHIP_NAME.lower() == "esp8266":
    url = 'https://raw.githubusercontent.com/uraich/uPyCraft-Qt5/master/microPythonFirmware/microPython-ESP8266.bin'
    sha256Hash="73f0b04e3356a6803bbcd7a32f3259ed2e32505601e3cd45628a21179f5a2fc4"
else:
    print("We don't have a micro Python version for the %s"%progFW.CHIP_NAME.lower())
    
savePath='firmware.bin'

with urllib.request.urlopen(url) as response, open(savePath, 'wb') as out_file:
    shutil.copyfileobj(response, out_file)

# calculate the hash code
sha256 = hashlib.sha256()

with open('firmware.bin', 'rb') as f:
    while True:
        data = f.read(BUF_SIZE)
        if not data:
            break
        sha256.update(data)

hashString = sha256.hexdigest()
if hashString != sha256Hash:
    print("Hash code of the downloaded firmware is wrong. Giving up...")
    sys.exit()
    
print("Hash code is correct! Erasing flash...")
print('---- erasing flash ----')
proc=subprocess.Popen(['esptool', 'erase_flash'],stdout=subprocess.PIPE)
for line in proc.stdout:
    lineString = line.strip().decode("utf-8")
    print(lineString)
sleep(0.5)
print('---- flash programming the %s with Micro Python ----'%progFW.CHIP_NAME.lower())

if progFW.CHIP_NAME.lower() == 'esp32':
    print('Calling the programming sub process')
    cmdLine='esptool write_flash --flash_size=detect --flash_mode dio 0x1000 firmware.bin'
else :
    cmdLine='esptool write_flash --flash_size=detect --flash_mode dio 0 firmware.bin'
    
args=shlex.split(cmdLine)
proc=subprocess.Popen(args,stdout=subprocess.PIPE)

for chunk in each_chunk(proc.stdout,'\r'):
    print(chunk+'\r',end='')