From df5ada68986c67750ed005542a2acee54f88edea Mon Sep 17 00:00:00 2001
From: Geoff Hutchison <geoff.hutchison@gmail.com>
Date: Wed, 6 Mar 2024 17:25:53 -0500
Subject: [PATCH] Add more error checking for Fetch PDB

Should fix #1616

Signed-off-by: Geoff Hutchison <geoff.hutchison@gmail.com>
---
 avogadro/qtplugins/fetchpdb/fetchpdb.cpp | 52 +++++++++++++++++-------
 1 file changed, 37 insertions(+), 15 deletions(-)

diff --git a/avogadro/qtplugins/fetchpdb/fetchpdb.cpp b/avogadro/qtplugins/fetchpdb/fetchpdb.cpp
index dbb5fd5f28..8b43c92e17 100644
--- a/avogadro/qtplugins/fetchpdb/fetchpdb.cpp
+++ b/avogadro/qtplugins/fetchpdb/fetchpdb.cpp
@@ -66,14 +66,6 @@ bool FetchPDB::readMolecule(QtGui::Molecule& mol)
 
 void FetchPDB::showDialog()
 {
-  if (!m_network) {
-    m_network = new QNetworkAccessManager(this);
-    connect(m_network, SIGNAL(finished(QNetworkReply*)), this,
-            SLOT(replyFinished(QNetworkReply*)));
-  }
-  if (!m_progressDialog) {
-    m_progressDialog = new QProgressDialog(qobject_cast<QWidget*>(parent()));
-  }
   // Prompt for a chemical structure name
   bool ok;
   QString pdbCode = QInputDialog::getText(
@@ -83,10 +75,36 @@ void FetchPDB::showDialog()
   if (!ok || pdbCode.isEmpty())
     return;
 
+  // check if the PDB code matches the expected format
+  if (pdbCode.length() != 4) {
+    QMessageBox::warning(qobject_cast<QWidget*>(parent()),
+                         tr("Invalid PDB Code"),
+                         tr("The PDB code must be exactly 4 characters long."));
+    return;
+  }
+
+  // first character should be 1-9
+  if (!pdbCode.at(0).isDigit() || pdbCode.at(0).toLatin1() == '0') {
+    QMessageBox::warning(
+      qobject_cast<QWidget*>(parent()), tr("Invalid PDB Code"),
+      tr("The first character of the PDB code must be 1-9."));
+    return;
+  }
+
+  if (!m_network) {
+    m_network = new QNetworkAccessManager(this);
+    connect(m_network, SIGNAL(finished(QNetworkReply*)), this,
+            SLOT(replyFinished(QNetworkReply*)));
+  }
+
   // Hard coding the PDB download URL
   m_network->get(QNetworkRequest(
     QUrl("https://files.rcsb.org/download/" + pdbCode + ".pdb")));
 
+  if (!m_progressDialog) {
+    m_progressDialog = new QProgressDialog(qobject_cast<QWidget*>(parent()));
+  }
+
   m_moleculeName = pdbCode;
   m_progressDialog->setLabelText(tr("Querying for %1").arg(pdbCode));
   m_progressDialog->setRange(0, 0);
@@ -106,15 +124,10 @@ void FetchPDB::replyFinished(QNetworkReply* reply)
   }
 
   m_moleculeData = reply->readAll();
-  m_tempFileName =
-    QDir::tempPath() + QDir::separator() + m_moleculeName + ".pdb";
-  QFile out(m_tempFileName);
-  out.open(QIODevice::WriteOnly);
-  out.write(m_moleculeData);
-  out.close();
 
   // Check if the file was successfully downloaded
-  if (m_moleculeData.contains("Error report") ||
+  if (m_moleculeData.contains("Not Found") ||
+      m_moleculeData.contains("Error report") ||
       m_moleculeData.contains("Page not found (404)")) {
     QMessageBox::warning(
       qobject_cast<QWidget*>(parent()), tr("Network Download Failed"),
@@ -122,7 +135,16 @@ void FetchPDB::replyFinished(QNetworkReply* reply)
     reply->deleteLater();
     return;
   }
+
+  m_tempFileName =
+    QDir::tempPath() + QDir::separator() + m_moleculeName + ".pdb";
+  QFile out(m_tempFileName);
+  out.open(QIODevice::WriteOnly);
+  out.write(m_moleculeData);
+  out.close();
+
   emit moleculeReady(1);
   reply->deleteLater();
 }
+
 } // namespace Avogadro::QtPlugins