4 files changed, 34 insertions, 1 deletions

diff --git a/mainwindow.cpp b/mainwindow.cpp
index a4e7bd1..41c4254 100644
--- a/mainwindow.cpp
+++ b/mainwindow.cpp
@@ -224,7 +224,9 @@ MainWindow::MainWindow(QWidget *parent)
             return;
 
         QString selectedFilter;
-        QString fileName = QFileDialog::getSaveFileName(this, tr("Save PCAP File"), "", tr("PCAP Files (*.pcap);;All Files (*.*)"), &selectedFilter);
+        QFileDialog saveDialog;
+        saveDialog.setDefaultSuffix("pcap");
+        QString fileName = saveDialog.getSaveFileName(this, tr("Save PCAP File"), "", tr("PCAP Files (*.pcap);;All Files (*.*)"), &selectedFilter);
         if (fileName.length() > 0) {
             pcpp::PcapFileWriterDevice pcapWriter(fileName.toStdString(), ppp->getLinkLayer());
             if (!pcapWriter.open())
@@ -303,6 +305,8 @@ MainWindow::MainWindow(QWidget *parent)
         }
         rawPacket->removeData(cursorPos, 1);
         rawPacket->insertData(cursorPos, reinterpret_cast<const uint8_t *>(cursorData.data()), cursorData.size());
+        const auto& row = currentSelectedRow();
+        ppp->fixupHeaders(row);
     });
 }
 
@@ -327,6 +331,15 @@ pcpp::RawPacket* MainWindow::currentSelectedPacket()
     return &ppp->getRawPacket(selected.last()->row());
 }
 
+ssize_t MainWindow::currentSelectedRow()
+{
+    const auto &selected = ui->tableWidget->selectedItems();
+    if (selected.empty())
+        return -1;
+
+    return selected.last()->row();
+}
+
 void MainWindow::closeEvent(QCloseEvent *closeEvent)
 {
     isAboutToClose = true;
diff --git a/mainwindow.h b/mainwindow.h
index 3e0a4a1..6c54370 100644
--- a/mainwindow.h
+++ b/mainwindow.h
@@ -22,6 +22,7 @@ public:
     MainWindow(QWidget *parent = nullptr);
     ~MainWindow();
     pcpp::RawPacket* currentSelectedPacket();
+    ssize_t currentSelectedRow();
     bool addTableRow();
     bool updateTableRow(size_t index);
 
diff --git a/pcapplusplus.cpp b/pcapplusplus.cpp
index 0848401..809c9e8 100644
--- a/pcapplusplus.cpp
+++ b/pcapplusplus.cpp
@@ -306,6 +306,24 @@ bool PcapPlusPlus::setPort(size_t index, const std::string & port, bool isSource
     return retval;
 }
 
+void PcapPlusPlus::fixupHeaders(size_t index)
+{
+    auto parsedPacket = pcpp::Packet(&getRawPacket(index));
+    auto * ip4Layer = parsedPacket.getLayerOfType<pcpp::IPv4Layer>();
+    auto * ip6Layer = parsedPacket.getLayerOfType<pcpp::IPv6Layer>();
+    auto * udpLayer = parsedPacket.getLayerOfType<pcpp::UdpLayer>();
+    auto * tcpLayer = parsedPacket.getLayerOfType<pcpp::TcpLayer>();
+
+    if (ip4Layer)
+        ip4Layer->computeCalculateFields();
+    if (ip6Layer)
+        ip6Layer->computeCalculateFields();
+    if (udpLayer)
+        udpLayer->computeCalculateFields();
+    if (tcpLayer)
+        tcpLayer->computeCalculateFields();
+}
+
 bool PcapPlusPlus::getPcapStatistics(pcpp::IFileDevice::PcapStats & stats)
 {
     if (reader) {
diff --git a/pcapplusplus.h b/pcapplusplus.h
index c0f0eb6..0b3e81a 100644
--- a/pcapplusplus.h
+++ b/pcapplusplus.h
@@ -27,6 +27,7 @@ public:
     bool setEth(size_t index, const std::string & eth, bool isSourceIp);
     bool setIp(size_t index, const std::string & ip, bool isSourceIp);
     bool setPort(size_t index, const std::string & port, bool isSourceIp);
+    void fixupHeaders(size_t index);
     bool getPcapStatistics(pcpp::IFileDevice::PcapStats & stats);
     double getFirstPacketTimestamp() { return firstPacketTs; }