From c9847645bab7e1a4293d95b18274411b7bd69c6d Mon Sep 17 00:00:00 2001
From: Torsten Harenberg <torsten@dl1thm.de>
Date: Tue, 21 Dec 2021 13:16:55 +0100
Subject: [PATCH] firmware extension checks implemented

---
 update/update.go | 49 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 47 insertions(+), 2 deletions(-)

diff --git a/update/update.go b/update/update.go
index 8939c43..a6a77c3 100644
--- a/update/update.go
+++ b/update/update.go
@@ -10,6 +10,7 @@ import (
 	"github.com/schollz/progressbar/v3"
 	"io/ioutil"
 	"os"
+	"path"
 	"regexp"
 	"strings"
 	"time"
@@ -104,9 +105,9 @@ func (m *modem) writeUpdate(fwfile string) error {
 		return err
 	}
 
-	log.Infof("flash stamp: % X", binary.LittleEndian.Uint32(answer[2:6]))
+	log.Debugf("flash stamp: % X", binary.LittleEndian.Uint32(answer[2:6]))
 	flashId := binary.LittleEndian.Uint16(answer[0:2])
-	log.Infof("flash id: % X", flashId)
+	log.Debugf("flash id: % X", flashId)
 	if flashId != 0xa41f && flashId != 0x5b1f && flashId != 0xda1f {
 		_, _ = m.write("\033", true) // send ESC to cancel the update
 		return errors.New("wrong flash id")
@@ -173,6 +174,47 @@ func (m *modem) writeUpdate(fwfile string) error {
 	return nil
 }
 
+func (m *modem) checkFirmwareFits(fwfile string) error {
+	answer, err := m.write("ver ##\r")
+	if err != nil {
+		log.Errorf("Error while reading the modem version: %s", err)
+		return err
+	}
+	re := regexp.MustCompile(`\w#1`)
+	version := strings.ReplaceAll(string(re.Find(answer)), "#1", "")
+	type modemtype struct {
+		ext, modem string
+	}
+	var extensions = map[string]modemtype{
+		"A": {".pt2", "PTC-II"},
+		"B": {".pro", "PTC-IIpro"},
+		"C": {".pte", "PTC-IIe"},
+		"D": {".pex", "PTC-IIex"},
+		"E": {".ptu", "PTC-IIusb"},
+		"F": {".ptn", "PTC-IInet"},
+		"H": {".dr7", "DR-7800"},
+		"I": {".dr7", "DR-7400"},
+		"K": {".pr7", "DR-7000"},
+		"L": {".p3u", "PTC-IIIusb"},
+		"T": {".ptx", "PTC-IIItrx"},
+	}
+
+	modem, exists := extensions[version]
+	if exists {
+		if modem.ext == path.Ext(fwfile) {
+			log.Infof("Firmware %s matches detected modem type %s.", fwfile, modem.modem)
+			return nil
+		} else {
+			log.Fatalf("Firmware %s does NOT match detected modem type %s.", fwfile, modem.modem)
+			return errors.New("Firmware and modem type do not match")
+		}
+	} else {
+		log.Fatal("Cannot detect modem type.")
+		return errors.New("Cannot detect modem type")
+	}
+
+}
+
 func Update(serport string, baudrate int, fwfile string) error {
 	m, err := initmodem(serport, baudrate)
 	defer m.port.Close()
@@ -190,6 +232,9 @@ func Update(serport string, baudrate int, fwfile string) error {
 	if m.setDateTime() != nil {
 		return err
 	}
+	if m.checkFirmwareFits(fwfile) != nil {
+		return err
+	}
 	if m.writeUpdate(fwfile) != nil {
 		return err
 	}