Update extract.py

extract.py

@@ -38,28 +38,48 @@ def get_vehicle_info(plate_number: str):
     driver = webdriver.Chrome(options=options)
 
     try:
-        # Step 1: Open the site
+        
         driver.get("https://www.jambisamsat.net/infopkb.html")
-        time.sleep(2)
+        time.sleep(1)
+
+        
+        # Wait until input box is present
+        WebDriverWait(driver, 10).until(
+            EC.presence_of_element_located((By.ID, "no_polisi"))
+        )
 
-        # Step 2: Find the input and enter plate number
-        input_element = driver.find_element(By.NAME, "nopol")
-        input_element.send_keys(plate_number)
+        # Fill in the plate number
+        input_field = driver.find_element(By.ID, "no_polisi")
+        input_field.clear()
+        input_field.send_keys(plate_number)
 
-        # Step 3: Submit the form
-        submit_button = driver.find_element(By.CSS_SELECTOR, 'input[type="submit"]')
+        # Click the submit button by class name
+        submit_button = driver.find_element(By.CSS_SELECTOR, 'button.btn.btn-primary[type="submit"]')
         submit_button.click()
-        time.sleep(2)
 
-        driver.implicitly_wait(5)
+        # Wait for the new page to load
+        WebDriverWait(driver, 10).until(
+            EC.url_contains("infopkb.php")
+        )
+
+        # # Step 2: Find the input and enter plate number
+        # input_element = driver.find_element(By.NAME, "nopol")
+        # input_element.send_keys(plate_number)
+
+        # # Step 3: Submit the form
+        # submit_button = driver.find_element(By.CSS_SELECTOR, 'input[type="submit"]')
+        # submit_button.click()
+        # time.sleep(2)
+
+        driver.implicitly_wait(3)
         
         page_title = driver.title
         screenshot = driver.get_screenshot_as_png()
 
         return Image.open(BytesIO(screenshot)) , page_title
 
-    except Exception as e:
-        return f"Error occurred: {e}"
+    except WebDriverException as e:
+        return Image.new('RGB', (1, 1)), page_title
 
     finally:
         driver.quit()