selview.py

import os
import argparse
import requests
import base64
import re
from requests.packages.urllib3.exceptions import InsecureRequestWarning
requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

http_proxy = ""
os.environ['HTTP_PROXY'] = http_proxy
os.environ['HTTPS_PROXY'] = http_proxy

def extract_ip(url):
    ip_pattern = r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}'
    match = re.search(ip_pattern, url)
    if match:
        return match.group()
    else:
        return None

def test(host_url, url_to_pass, browser):
    # Make the first POST request to get the session ID
    response = requests.post(f'{host_url}/wd/hub/session', 
                         json={
                             "desiredCapabilities": {
                                 "browserName": browser,
                                 "version": "",
                                 "platform": "ANY",
                                 "enableVNC": True,
                                 "name": "screenshot-test",
                                 "sessionTimeout": "120s"
                             },
                             "args": ["--headers=Metadata-Flavor: Google"]
                         },
                         verify=False,
                         timeout=60)

    
    if response.status_code != 200:
        print("Got an error from the server: " + response.text + " and HTTP code: " + str(response.status_code))
        exit(0)
    
    if "create container: Error response from daemon" in response.text:
        print("Host has an issue: " + response.text)
        exit(0)
    
    try:
        session_id = response.json()['sessionId']
    except:
        session_id = response.json()['value']['sessionId']
    
    if session_id is None:
        match = re.search(r'"sessionId":"([^"]+)"', response.text)
        if match:
            session_id = match.group(1)
            print(session_id)
        else:
            print("sessionId not found")
            print(response.text)

    # Make the second POST request to set the URL
    headers = {"Metadata-Flavor": "Google"}
    response = requests.post(f'{host_url}/wd/hub/session/{session_id}/url', verify=False, json={"url": f"{url_to_pass}"},headers=headers)

    # Make the GET request to get the screenshot
    response = requests.get(f'{host_url}/wd/hub/session/{session_id}/screenshot', verify=False)
    screenshot_data = response.json()['value']
    # take the ip and extract it.
    ip = extract_ip(host_url)
    # Save the screenshot to a file
    with open('./output/'+ip+'.png', 'wb') as file:
        file.write(base64.b64decode(screenshot_data))
    
    #os.system('open ./output/'+ip+'.png')
if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("-s", "--server", required=True, help="URL of Selenoid UI")
   parser.add_argument("-u", "--url", required=True, help="URL to get screenshot of use file:// for files.")
   parser.add_argument("-b", "--browser", required=False, default="chrome" ,help="Try a different browser. firefox chrome safari opera")
   args = parser.parse_args()
   host_url = args.server
   url_to_pass = args.url
   if host_url[-1] == "/":
      host_url = host_url[:-1]
   browser = args.browser
   test(host_url,url_to_pass,browser)