Merge branch 'main' into zdy

This commit is contained in:
David Chang
2024-02-15 22:27:10 +08:00
153 changed files with 5317 additions and 740 deletions

View File

@@ -1,24 +1,29 @@
import json
import logging
import os
import os.path
import sqlite3
import tempfile
import time
import traceback
import uuid
import tempfile
from datetime import datetime, timedelta
from typing import Any, Union, Optional
from typing import Dict, List
import os
import shutil
import requests
from playwright.sync_api import sync_playwright, TimeoutError
from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive, GoogleDriveFile, GoogleDriveFileList
from playwright.sync_api import sync_playwright, TimeoutError
from requests_toolbelt.multipart.encoder import MultipartEncoder
from desktop_env.controllers.python import PythonController
from desktop_env.evaluators.metrics.utils import compare_urls
logger = logging.getLogger("desktopenv.setup")
FILE_PATH = os.path.dirname(os.path.abspath(__file__))
class SetupController:
def __init__(self, vm_ip: str, cache_dir: str):
@@ -130,7 +135,8 @@ class SetupController:
break
except requests.RequestException as e:
logger.error(f"Failed to download {url} caused by {e}. Retrying... ({max_retries - i - 1} attempts left)")
logger.error(
f"Failed to download {url} caused by {e}. Retrying... ({max_retries - i - 1} attempts left)")
if not downloaded:
raise requests.RequestException(f"Failed to download {url}. No retries left. Error: {e}")
@@ -349,18 +355,18 @@ class SetupController:
logger.info("Connect to Chrome @: %s", remote_debugging_url)
logger.debug("PLAYWRIGHT ENV: %s", repr(os.environ))
for attempt in range(15):
if attempt>0:
if attempt > 0:
time.sleep(5)
browser = None
with sync_playwright() as p:
try:
browser = p.chromium.connect_over_cdp(remote_debugging_url)
#break
# break
except Exception as e:
if attempt < 14:
logger.error(f"Attempt {attempt + 1}: Failed to connect, retrying. Error: {e}")
#time.sleep(10)
# time.sleep(10)
continue
else:
logger.error(f"Failed to connect after multiple attempts: {e}")
@@ -379,7 +385,7 @@ class SetupController:
try:
page.goto(url, timeout=60000)
except:
logger.warning("Opening %s exceeds time limit", url) # only for human test
logger.warning("Opening %s exceeds time limit", url) # only for human test
logger.info(f"Opened tab {i + 1}: {url}")
if i == 0:
@@ -458,16 +464,17 @@ class SetupController:
for p in paths:
q = f'"{parent_id}" in parents and title = "{p}" and mimeType = "application/vnd.google-apps.folder" and trashed = false'
folder = drive.ListFile({'q': q}).GetList()
if len(folder) == 0: # not exists, create it
if len(folder) == 0: # not exists, create it
parents = {} if parent_id == 'root' else {'parents': [{'id': parent_id}]}
file = drive.CreateFile({'title': p, 'mimeType':'application/vnd.google-apps.folder', **parents})
file = drive.CreateFile({'title': p, 'mimeType': 'application/vnd.google-apps.folder', **parents})
file.Upload()
parent_id = file['id']
else: parent_id = folder[0]['id']
else:
parent_id = folder[0]['id']
return parent_id
for oid, operation in enumerate(config['operation']):
if operation == 'delete': # delete a specific file
if operation == 'delete': # delete a specific file
# query pattern string, by default, remove all files/folders not in the trash to the trash
params = config['args'][oid]
q = params.get('query', '')
@@ -476,15 +483,19 @@ class SetupController:
filelist: GoogleDriveFileList = drive.ListFile({'q': q_file}).GetList()
q_folder = f"( {q} ) and mimeType = 'application/vnd.google-apps.folder'" if q.strip() else "mimeType = 'application/vnd.google-apps.folder'"
folderlist: GoogleDriveFileList = drive.ListFile({'q': q_folder}).GetList()
for file in filelist: # first delete file, then folder
for file in filelist: # first delete file, then folder
file: GoogleDriveFile
if trash: file.Trash()
else: file.Delete()
if trash:
file.Trash()
else:
file.Delete()
for folder in folderlist:
folder: GoogleDriveFile
# note that, if a folder is trashed/deleted, all files and folders in it will be trashed/deleted
if trash: folder.Trash()
else: folder.Delete()
if trash:
folder.Trash()
else:
folder.Delete()
elif operation == 'mkdirs':
params = config['args'][oid]
mkdir_in_googledrive(params['path'])
@@ -508,7 +519,6 @@ class SetupController:
else:
raise ValueError('[ERROR]: not implemented clean type!')
def _login_setup(self, **config):
""" Login to a website with account and password information.
@args:
@@ -568,3 +578,82 @@ class SetupController:
raise NotImplementedError
return browser, context
def _update_browse_history_setup(self, **config):
db_path = os.path.join("desktop_env", "assets", "history_empty.sqlite")
# copy a new history file in the tmp folder
cache_path = os.path.join(self.cache_dir, "history_new.sqlite")
shutil.copyfile(db_path, cache_path)
db_path = cache_path
history = config['history']
for history_item in history:
url = history_item['url']
title = history_item['title']
visit_time = datetime.now() - timedelta(seconds=history_item['visit_time_from_now_in_seconds'])
# Chrome use ms from 1601-01-01 as timestamp
epoch_start = datetime(1601, 1, 1)
chrome_timestamp = int((visit_time - epoch_start).total_seconds() * 1000000)
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute('''
INSERT INTO urls (url, title, visit_count, typed_count, last_visit_time, hidden)
VALUES (?, ?, ?, ?, ?, ?)
''', (url, title, 1, 0, chrome_timestamp, 0))
url_id = cursor.lastrowid
cursor.execute('''
INSERT INTO visits (url, visit_time, from_visit, transition, segment_id, visit_duration)
VALUES (?, ?, ?, ?, ?, ?)
''', (url_id, chrome_timestamp, 0, 805306368, 0, 0))
conn.commit()
conn.close()
logger.info('Fake browsing history added successfully.')
controller = PythonController(self.vm_ip)
# get the path of the history file according to the platform
os_type = controller.get_vm_platform()
if os_type == 'Windows':
chrome_history_path = controller.execute_python_command(
"""import os; print(os.path.join(os.getenv('USERPROFILE'), "AppData", "Local", "Google", "Chrome", "User Data", "Default", "History"))""")[
'output'].strip()
elif os_type == 'Darwin':
chrome_history_path = controller.execute_python_command(
"""import os; print(os.path.join(os.getenv('HOME'), "Library", "Application Support", "Google", "Chrome", "Default", "History"))""")[
'output'].strip()
elif os_type == 'Linux':
chrome_history_path = controller.execute_python_command(
"import os; print(os.path.join(os.getenv('HOME'), '.config', 'google-chrome', 'Default', 'History'))")[
'output'].strip()
else:
raise Exception('Unsupported operating system')
form = MultipartEncoder({
"file_path": chrome_history_path,
"file_data": (os.path.basename(chrome_history_path), open(db_path, "rb"))
})
headers = {"Content-Type": form.content_type}
logger.debug(form.content_type)
# send request to server to upload file
try:
logger.debug("REQUEST ADDRESS: %s", self.http_server + "/setup" + "/upload")
response = requests.post(self.http_server + "/setup" + "/upload", headers=headers, data=form)
if response.status_code == 200:
logger.info("Command executed successfully: %s", response.text)
else:
logger.error("Failed to upload file. Status code: %s", response.text)
except requests.exceptions.RequestException as e:
logger.error("An error occurred while trying to send the request: %s", e)
self._execute_setup(["sudo chown -R user:user /home/user/.config/google-chrome/Default/History"], shell=True)