app.py

#from __future__ import annotations
from selenium import webdriver
import gradio as gr
import uuid
import re
from PIL import Image
from io import BytesIO
from pathlib import Path
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
#from selenium_webdriver import WebElement

driver_type = 'chromedriver'
driver=False

def get_concat_h_cut(in1, in2):
    im1=Image.open(in1)
    im2=Image.open(in2)
    dst = Image.new('RGB', (im1.width + im2.width,
                            min(im1.height, im2.height)))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (im1.width, 0))
    return dst


def get_concat_v_cut(in1, in2):
    im1=Image.open(in1)
    im2=Image.open(in2)    
    dst = Image.new(
        'RGB', (min(im1.width, im2.width), im1.height + im2.height))
    dst.paste(im1, (0, 10))
    dst.paste(im2, (0, im1.height+10))
    return dst

css_="""--solid_color_primary: #132241;--text_color_primary: #ffffff;--solid_color_bubble: #2c3d7d;--solid_color_tertiary: #223e69;--solid_color_block: #1a1f23;"""
css_bub="""padding: 20px!important;"""
def run_script(url: str, height: int, width: int, check_b,check_h):
    mes_box=[]
    out_box=[]
    uid=uuid.uuid4()
    out=None
    is_url=True
    if is_url:
        options = webdriver.ChromeOptions()
        options.add_argument('--headless')
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')
        mes='<center>operation success'
        try:
            driver = webdriver.Chrome(options=options)
            driver.get(url)
            html=driver.page_source
            #print (driver.html)
            driver.implicitly_wait(30)

            

            driver.set_window_size(int(width), int(height))
            page = driver.find_element(By.TAG_NAME, "html")
            driver.execute_script(f"arguments[0].style='{css_}';",page)
            print (dir(page))
            #page.send_keys('theme','gab-black')
            dk=page.get_attribute('theme')
            print (dir(dk))
            print (page.get_attribute('theme'))
            #driver.execute_script("arguments[0].theme='gab-black';",page)
            print (page.get_attribute('theme'))
            print(dir(driver))
            #driver.get(driver.getCurrentUrl())
            #driver.navigate().refresh()
            main_head = driver.find_element(By.CLASS_NAME, "main-content-header")
            #driver.execute_script("arguments[0].style='background:#474747;color:#055550!important;';",main_head)
            
            head_shot = main_head.screenshot(f'head-{uid}-tmp.png')
            
            obj = driver.find_element(By.CLASS_NAME, "main")            
            messages = driver.find_elements(By.CLASS_NAME, "message")
            message_css = driver.find_element(By.CSS_SELECTOR, ".message-bubble")
            driver.execute_script(f"arguments[0].style='{css_bub}';",message_css)

            #print (len(messages))
            #mes_blocks=len(messages)%2
            
            for i,ea in enumerate(check_b):
                try:    
                    ea = int(ea)
                    ea = (ea*2)-2
                    print (ea)
                    messages[ea].screenshot(f'{i}-{uid}-tmp.png')
                    messages[ea+1].screenshot(f'{i+1}-{uid}-tmp.png')
                    im_roll = get_concat_v_cut(f'{i}-{uid}-tmp.png',f'{i+1}-{uid}-tmp.png')
                    im_roll.save(f'comb{i}-{uid}-tmp.png')
                    out_box.append(f'comb{i}-{uid}-tmp.png')
                

                    print(f'out_box:: {out_box}')
                    #mes_box.clear()
                except Exception:
                    mes="<center>Some blocks returned an error"
            if out_box:
                if len(out_box)>1:
                    im_roll = get_concat_v_cut(f'{out_box[0]}',f'{out_box[1]}')
                    im_roll.save(f'comb-{uid}-tmp.png')                
                    for i in range(2,len(out_box)):
                        im_roll = get_concat_v_cut(f'comb-{uid}-tmp.png',f'{out_box[i]}')
                        im_roll.save(f'comb-{uid}-tmp.png')
                    out = f'comb-{uid}-tmp.png'
                else:
                    tmp_im = Image.open(out_box[0])
                    tmp_im.save(f'comb-{uid}-tmp.png')
                    out = f'comb-{uid}-tmp.png'
                if check_h:
                    head_paste = get_concat_v_cut(f'head-{uid}-tmp.png',f'comb-{uid}-tmp.png')
                    head_paste.save(f'comb-{uid}-tmp.png')
                    out = f'comb-{uid}-tmp.png'
                        
            
            
            
            
            '''
            for i,mes in enumerate(messages):
                print (mes.text)
                if not (i%2):
                    messages[i].screenshot(f'{i}-{uid}-tmp.png')
                    messages[i+1].screenshot(f'{i+1}-{uid}-tmp.png')
                    mes_box.append(())
                    out_box.append(get_concat_v_cut(f'{i}-{uid}-tmp.png',f'{i+1}-{uid}-tmp.png'))
                    print(f'out_box:: {out_box}')
                    #mes_box.clear()
            '''    
            #obj = driver.find_element(By.ID, "app-container")            
            #obj = driver.find_element(By.ID, "conversation-feed")            
            #obj = driver.find_element_by_id("in_html")
            ######################
            #inp = driver.find_element("id","input")
            '''
            find_element(By.ID, "id")
            find_element(By.NAME, "name")
            find_element(By.XPATH, "xpath")
            find_element(By.LINK_TEXT, "link text")
            find_element(By.PARTIAL_LINK_TEXT, "partial link text")
            find_element(By.TAG_NAME, "tag name")
            find_element(By.CLASS_NAME, "class name")
            find_element(By.CSS_SELECTOR, "css selector")       
            '''
            #inp.send_keys("test")
            ######################
            cookie_jar = []
            #cookie_jar.append(driver.get_cookies())
            #print(cookie_jar)
            screenshot = obj.screenshot(f'{uid}-tmp.png')
            #screenshot = obj.get_screenshot_as_png()
        except WebDriverException as e:
            return Image.new('RGB', (1, 1)), f'<center>{e}',out_box,out
        finally:
            if driver:
                driver.quit()

        #return [Image.open(BytesIO(screenshot)), 'operation success.',cookie_jar,html]
        return Image.open(f'{uid}-tmp.png'), mes,out_box,out
    else:
        return None, '<center>Please enter a valid URL of a website/host.',out_box,out

with gr.Blocks() as app:
    
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="URL",lines=1)
            btn= gr.Button()
            
        with gr.Column():
            check_h=gr.Checkbox(label="Show Header", value=True)
            check_b=gr.CheckboxGroup(label="Chatblocks", choices=[1,2,3,4,5,6,7,8,9,10], value=[1])
            with gr.Row():
                height=gr.Number(label="Height", value=4096)
                width=gr.Number(label="Width",value=800)
    message=gr.HTML('<center>Enter URL')
    with gr.Row():
        with gr.Column():
            out=gr.Image()
            outgal=gr.Gallery()
        outim = gr.Image()

    btn.click(run_script,[inp,height,width,check_b,check_h],[outim,message,outgal,out])
app.launch()