ESP32Flash.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Net.Http.Headers;
using System.Security.Cryptography;
using System.Security.Policy;
using System.Text;

namespace ESP32
{
    public enum ESP32MemoryRegion
    {
        IROM0, IRAM0, IRAM1, DROM0, DRAM0, DRAM1
    }

    public class ESP32MemoryMap
    {
        internal struct MemoryRange
        {

            internal MemoryRange(long start, long size)
            {
                Start = start;
                Size = size;
            }
            internal long Start;
            internal long Size;

            internal long Stop
            {
                get
                {
                    return Start + Size;
                }
            }
        }

        static Dictionary<MemoryRange, ESP32MemoryRegion> Map = new Dictionary<MemoryRange, ESP32MemoryRegion>();

        static ESP32MemoryMap()
        {
            Map.Add(new MemoryRange(0x40800000, 4194304), ESP32MemoryRegion.IROM0);
            Map.Add(new MemoryRange(0x40000000, 4194304), ESP32MemoryRegion.IRAM0);
            Map.Add(new MemoryRange(0x40400000, 4194304), ESP32MemoryRegion.IRAM1);
            Map.Add(new MemoryRange(0x3F400000, 4194304), ESP32MemoryRegion.DROM0);
            Map.Add(new MemoryRange(0x3FF80000, 524288), ESP32MemoryRegion.DRAM0);
            Map.Add(new MemoryRange(0x3F800000, 4194304), ESP32MemoryRegion.DRAM1);
        }

        internal static ESP32MemoryRegion GetRegionForAddr(long addr)
        {
            return Map.Where(r => addr >= r.Key.Start && addr < r.Key.Stop).FirstOrDefault().Value;
        }
    }

    public class ESP32Flash
    {

        public string SourceFile = "";
        public ESP32AppImage BootLoaderImage;
        public List<ESP32Partition> Partitions = new List<ESP32Partition>();

        public void WriteToFile(string filePath)
        {
            BinaryWriter bw = new BinaryWriter(File.OpenWrite(filePath));
            
            /* first 0x1000 are 0xFF bytes */
            for(var x =0; x < 0x1000; x++)
            {
                bw.Write((byte)0xFF);
            }

            /* next up is Bootloader Image */
            bw.Write(BootLoaderImage.GetBytes());

            /* fill remaining bootloader space */
            while (bw.BaseStream.Position < 0x8000)
            {
                bw.Write((byte)0xFF);
            }

            foreach(var partition in Partitions)
            {
                /* Partition Record Header */
                bw.Write((ushort)0x50AA);

                bw.Write(partition.Type);
                bw.Write(partition.SubType);
                bw.Write(partition.Offset);
                bw.Write(partition.Data.Length);

                byte[] nameBuffer = new byte[20];
                byte[] nameBytes = UTF8Encoding.UTF8.GetBytes(partition.Name);
                Array.Copy(nameBytes, 0, nameBuffer, 0, nameBytes.Length);
                bw.Write(nameBuffer);
            }

            /* fill remaining partition table space */
            int padSize = (int)(0x9000 - bw.BaseStream.Position);
            byte[] pad = new byte[padSize];
            for (var x = 0; x < pad.Length; x++)
            {
                pad[x] = 0xFF;
            }
            bw.Write(pad);
            
            /* write the partition data */
            foreach (var partition in Partitions)
            {

                /* pad to next offset */
                padSize = (int)(partition.Offset - bw.BaseStream.Position);
                pad = new byte[padSize];
                for (var x = 0; x < pad.Length; x++)
                {
                    pad[x] = 0xFF;
                }
                bw.Write(pad);

                bw.Write(partition.Data);


            }

            bw.Close();
        }

        public void CorrectChecksums()
        {
            /* walk the app images and correct the checksums */
        }

        public ESP32Flash(string filePath)
        {
            SourceFile = filePath;
            BinaryReader br = new BinaryReader(new FileStream(filePath, FileMode.Open));

            /* BootROM is loaded to the first 0x1000 bytes, its all 0xFF in a flash dump, skip it */
            br.BaseStream.Position = 0x1000;


            /* Bootloader is up to 0x7000 bytes */
            BootLoaderImage = new ESP32AppImage(br.ReadBytes(0x7000));

            byte[] partitionRecord = br.ReadBytes(0x20);
            while (BitConverter.ToUInt16(partitionRecord, 0) == 0x50AA)
            {
                /* This is a partition! */
                ESP32Partition newPartition = new ESP32Partition();
                newPartition.Type = partitionRecord[0x02];
                newPartition.SubType = partitionRecord[0x03];
                newPartition.Offset = BitConverter.ToUInt32(partitionRecord, 0x04);
                newPartition.Length = BitConverter.ToUInt32(partitionRecord, 0x08);
                newPartition.Name = UTF8Encoding.UTF8.GetString(partitionRecord, 0x0C, 20).Split('\0')[0];

                var oldOffset = br.BaseStream.Position;
                br.BaseStream.Position = newPartition.Offset;
                newPartition.Data = br.ReadBytes((int)newPartition.Length);
                br.BaseStream.Position = oldOffset;

                Partitions.Add(newPartition);

                partitionRecord = br.ReadBytes(0x20);
            }
            br.Close();
        }

        public ESP32Partition GetPartitionByName(string name)
        {
            return Partitions.Where(p => p.Name.Equals(name)).FirstOrDefault();
        }

    }

    public struct ImagePatch
    {
        public int Offset;
        public byte[] Data;
    }

    public class ESP32Partition
    {
        public byte Type { get; internal set; }
        public byte SubType { get; internal set; }
        public string Name { get; internal set; }
        public byte[] Data { get; internal set; }
        public uint Offset { get; internal set; }
        public uint Length { get; internal set; }

        public ESP32AppImage ParseAppImage()
        {
            if (Type != 0x00)
            {
                return null;
            }

            if (Data[0] == 0xE9)
            {
                return new ESP32AppImage(Data);
            }
            else
            {
                return null;
            }
        }

        public void UpdateFromAppImage(ESP32AppImage img)
        {
            byte[] imgData = img.GetBytes();
            Length = (uint)imgData.Length;
            Data = imgData;
        }
    }

    public class ESP32AppImage
    {
        public byte Magic { get; internal set; }
        public byte SegmentCount { get; internal set; }
        public byte SPIMode { get; internal set; }
        public byte SPISpeed { get; internal set; }
        public byte SPISize { get; internal set; }
        public uint EntryAddress { get; internal set; }
        public byte WPPin { get; internal set; }
        public byte[] SPIPinDrv { get; internal set; }
        public ushort ChipID { get; internal set; }
        public byte MinChipRev { get; internal set; }
        public byte[] Reserved { get; internal set; }
        public byte HashAppended { get; internal set; }
        public List<ESP32AppSegment> Segments = new List<ESP32AppSegment>();

        public byte Checksum { get; internal set; }
        public byte[] ValidationHash;


        public byte[] CalcAppImageHash()
        {
            MemoryStream ms = new MemoryStream();
            BinaryWriter bw = new BinaryWriter(ms);

            bw.Write(this.Magic);
            bw.Write(this.SegmentCount);
            bw.Write(this.SPIMode);
            bw.Write((byte)((this.SPISize << 4) | this.SPISpeed));
            bw.Write(this.EntryAddress);

            bw.Write(this.WPPin);
            bw.Write(this.SPIPinDrv);
            bw.Write(this.ChipID);
            bw.Write(this.MinChipRev);
            bw.Write(this.Reserved);
            bw.Write(this.HashAppended);

            for (var x = 0; x < this.SegmentCount; x++)
            {
                ESP32AppSegment segment = this.Segments[x];
                bw.Write(segment.LoadAddress);
                bw.Write(segment.Length);
                bw.Write(segment.Data);
            }

            /* get to 16 byte boundary */
            while ((bw.BaseStream.Position + 1) % 0x10 != 0)
            {
                bw.Write((byte)0x00);
            }
            bw.Write(this.Checksum);

            var imageData = ms.ToArray();


            var sha256 = SHA256.Create();
            var calcHash = sha256.ComputeHash(imageData);
            int i = 3;
            return calcHash;
        }

        public bool IsChecksumValid()
        {

            return CalcChecksumByte() == this.Checksum;
            /*         for b in data:
            if type(b) is int:  # python 2/3 compat
                state ^= b
            else:
                state ^= ord(b)

        return state */
        }

        internal byte CalcChecksumByte()
        {
            byte state = 0xef;
            foreach (var seg in Segments)
            {
                foreach (var b in seg.Data)
                {
                    state ^= b;
                }
            }

            return state;
        }

        internal byte[] GetBytes()
        {
            MemoryStream ms = new MemoryStream();
            BinaryWriter bw = new BinaryWriter(ms);

            bw.Write(Magic);
            bw.Write((byte)Segments.Count);
            bw.Write(SPIMode);
            byte SPIByte = (byte)((SPISize << 4) | (SPISpeed & 0xF));
            bw.Write(SPIByte);

            bw.Write(EntryAddress);
            bw.Write(WPPin);
            bw.Write(SPIPinDrv);
            bw.Write(ChipID);
            bw.Write(MinChipRev);
            bw.Write(Reserved);
            bw.Write(HashAppended);

            foreach(var segment in Segments)
            {
                bw.Write(segment.LoadAddress);
                bw.Write(segment.Length);
                bw.Write(segment.Data);
            }

            while ((bw.BaseStream.Position + 1) % 0x10 != 0)
            {
                bw.Write((byte)0x00);
            }
            bw.Write(Checksum);
            if (HashAppended == 0x01)
            {
                bw.Write(ValidationHash);
            }

            byte[] imageData = ms.ToArray();
            bw.Close();
            return imageData;
        }

        public ESP32AppImage(byte[] Data)
        {
            MemoryStream ms = new MemoryStream(Data);
            BinaryReader br = new BinaryReader(ms);

            var magicCheck = br.ReadByte();
            if (magicCheck != 0xE9)
            {
                br.Close();
                throw new Exception("App Image Magic is not correct. Expected 0xE9");
            }

            this.Magic = magicCheck;
            this.SegmentCount = br.ReadByte();
            this.SPIMode = br.ReadByte();
            byte SPIByte = br.ReadByte();
            this.SPISpeed = (byte)(SPIByte & 0xF);
            this.SPISize = (byte)(SPIByte >> 4);
            this.EntryAddress = BitConverter.ToUInt32(br.ReadBytes(4),0);
            this.WPPin = br.ReadByte();
            this.SPIPinDrv = br.ReadBytes(3);
            this.ChipID = BitConverter.ToUInt16(br.ReadBytes(2),0);
            this.MinChipRev = br.ReadByte();
            this.Reserved = br.ReadBytes(8);
            this.HashAppended = br.ReadByte();

            for (var x = 0; x < this.SegmentCount; x++)
            {
                ESP32AppSegment segment = new ESP32AppSegment();
                segment.LoadAddress = BitConverter.ToUInt32(br.ReadBytes(4),0);
                segment.Length = BitConverter.ToUInt32(br.ReadBytes(4),0);
                segment.PhysicalOffset = (uint)br.BaseStream.Position;
                segment.Data = br.ReadBytes((int)segment.Length);
                this.Segments.Add(segment);
            }

            /* get to 16 byte boundary */
            while ((br.BaseStream.Position + 1) % 0x10 != 0)
            {
                br.ReadByte();
            }

            this.Checksum = br.ReadByte();
            if (this.HashAppended == 0x01)
            {
                this.ValidationHash = br.ReadBytes(0x20);
            }
            br.Close();
        }

    }

    public class ESP32AppSegment
    {
        public uint PhysicalOffset { get; internal set; }
        public uint LoadAddress { get; internal set; }
        public uint Length { get; internal set; }
        public byte[] Data { get; internal set; }

        public ESP32MemoryRegion Region
        {
            get
            {
                return ESP32MemoryMap.GetRegionForAddr(LoadAddress);
            }
        }

        public override string ToString()
        {
            return "0x" + LoadAddress.ToString("X") + " - " + Region.ToString();
        }
    }

}