Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import socket
 import requests
+import pycountry
 from urllib.parse import urlparse
 
 # Function to extract the domain name from the URL
@@ -32,13 +33,17 @@ def get_country_info(ip):
 
 # Function to get the flag URL from a country code
 def get_country_flag(country_name):
-    country_flags = {
-        "US": "https://countryflagsapi.com/png/us",
-        "IN": "https://countryflagsapi.com/png/in",
-        "GB": "https://countryflagsapi.com/png/gb",
-        # Add more country codes and flags as needed
-    }
-    return country_flags.get(country_name.upper(), "https://countryflagsapi.com/png/unknown")
+    try:
+        # Using pycountry to get the alpha_2 country code (e.g., "US", "IN")
+        country = pycountry.countries.get(name=country_name)
+        if country:
+            country_code = country.alpha_2
+            flag_url = f"https://countryflagsapi.com/png/{country_code.lower()}"
+            return flag_url
+        else:
+            return "https://countryflagsapi.com/png/unknown"  # Default flag if country not found
+    except AttributeError:
+        return "https://countryflagsapi.com/png/unknown"  # Return unknown flag if error
 
 # Streamlit app
 def main():