利用selenium实现截图

package com.dinfo.hrb.util;


import java.awt.image.BufferedImage;
import java.io.BufferedInputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;


import javax.imageio.ImageIO;
import javax.imageio.stream.FileImageOutputStream;
import javax.sql.rowset.serial.SerialBlob;


import org.apache.commons.io.FileUtils;
import org.openqa.selenium.Alert;
import org.openqa.selenium.Dimension;
import org.openqa.selenium.JavascriptExecutor;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.Point;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.WebElement;
import org.openqa.selenium.firefox.FirefoxDriver;


import com.dinfo.plugtool.common.AttributeCommon;
import com.dinfo.plugtool.dao.DBUtil;
import com.dinfo.plugtool.util.PropertyUtil;
import com.dinfo.plugtool.util.ScreenChecked;


/**
* 截图
* @author landmark
*
*/
public class ScreenShot implements Runnable {


private static SimpleDateFormat sdf3 = new SimpleDateFormat("yyyy-MM-dd");
private WebDriver driver;
private String windowHandle;
private int counts;
private String name;
private int flags = 0;
private String file_path;
private int unique_code;


public ScreenShot() {}
public ScreenShot(String name,String file_path,int unique_code) {
this.name = name;
this.file_path = file_path;
this.unique_code = unique_code;
}

public void run(){

while(true){
Map<String, Object> data = this.getData();
try {
start(data);
} catch (Exception e) {
try { driver.close(); } catch (Exception e1) {}
driver = null;
}
}
}

private Map<String,Object> getData(){

Map<String, Object> map = null;
try {
map = DatasHandle.getData();
} catch (Exception e2) {
e2.printStackTrace();
return null;
}
if(map == null){
System.out.println("未获取到数据，线程等待");
flags++;
if(flags == 20){
flags = 0;
try {
driver.close();
} catch (Exception e) {}
driver = null;
}
if(driver == null){
System.out.println("线程沉睡10分钟");
String sql_5 = "UPDATE r_sitemap_new SET DTL_ID='' WHERE ISTAKESCREEN IN(2,3,4,5,9) AND PROCESS_DATE='"+sdf3.format(new Date())+"'";
try {
DBUtil.update(sql_5);
} catch (SQLException e) {
e.printStackTrace();
}
try { Thread.sleep(1000*60*9); } catch (InterruptedException e) {}
}
else{
System.out.println("线程沉睡30秒");
try { Thread.sleep(1000*30); } catch (InterruptedException e) {}
}
return null;
}
return map;
}

@SuppressWarnings("static-access")
public void start(Map<String, Object> map){

//检测是否需要截图
if(!ScreenChecked.checked(map)){
return;
}

if(driver == null){
init();
}
else{
try {
driver.getWindowHandles();
} catch (Exception e) {
try {
driver.close();
} catch (Exception e1) {}
init();
}
}

String id = map.get("ID").toString();
String product_name = map.get("PRODUCT_NAME").toString();
String url = map.get("URL").toString();
String date = map.get("PROCESS_DATE").toString();
String bank_name = map.get("BANK_NAME").toString();

//打开需截图的页面
try {
driver.get(url);
} catch (Exception e2) {
driver.get("http://www.baidu.com/");
return;
}
//再次重新打开截图页面(针对南京银行页面url重定向问题)
if(bank_name.equals("南京银行")){
try {
driver.get(url);
} catch (Exception e2) {
driver.get("http://www.baidu.com/");
return;
}
}
if(bank_name.equals("农业银行") && (product_name.equals("微博"))){
try {
Thread.currentThread().sleep(6000);
} catch (InterruptedException e) {}
}
if(bank_name.equals("农业银行") &&
product_name.equals("泰康人寿保险股份有限公司")){
JavascriptExecutor js = (JavascriptExecutor)driver;
int i = 0;
int j = 100;
for (; i < 3000;) {
js.executeScript("scrollTo(" + i + "," + j + ")");
i += 100;
j += 100;
}
}

//检测页面是否是360主页
boolean c = ScreenChecked.checkUrlISMainPage(driver);
if(c){
System.out.println("页面检测不通过，不满足截图条件-----"+id);
DatasHandle.update_404(id);
driver.get("http://www.baidu.com/");
return;
}

//跳转弹出框，关闭弹出框
try {
Alert alert = driver.switchTo().alert();
alert.dismiss();
} catch (Exception e2) {}

//关闭多余弹出页
Set<String> set = null;
try {
set = driver.getWindowHandles();
} catch (Exception e1) {
flags++;
if(flags == 20){
flags = 0;
try { driver.close(); } catch (Exception e) {}
driver = null;
}
return;
}

try {
if (set != null && set.size() > 1) {

for (String string : set) {
if (string.equals(windowHandle)) {
continue;
} else {
driver.switchTo().window(string);
try {
Thread.sleep(500);
} catch (InterruptedException e) {
e.printStackTrace();
}
driver.close();
}
}
driver.switchTo().window(windowHandle);
}
} catch (Exception e1) {
e1.printStackTrace();
}

//页面检测：是否是404页面，或站内自定义错误页面
try {
PageCheck.checked(driver);
} catch (Exception e4) {
System.out.println("页面检测不通过，不满足截图条件-----"+id);
//检测是否属于站内页面异常
boolean b = PageCheck.checkURL(url,date);
if(b){
DatasHandle.update_itself(id);
}
else{
DatasHandle.update_404(id);
}
driver.get("http://www.baidu.com/");
return;
}

//截图
String path = "";
try {
byte[] b = takeScreenshot();
bytesToImageFile(b);
while (true) {
//获取图片文件大小
File file_ = new File(file_path);
long length = file_.length();
System.out.println(length);
//如果图片大于指定大小
int size = Integer.parseInt(PropertyUtil.getValue("max_size"));
if (length > size * 1024L) {
//进行图片压缩
path = ImageCompressUtil.zipImageFile(file_path, 0,
0, 0.8F, "_" + unique_code);
if(new File(path).length() > length){
break;
}
try {
FileUtils.copyFile(new File(path), new File(file_path));
} catch (Exception e) {
e.printStackTrace();
}
new File(path).delete();
} else {
break;
}
}

byte[] bytes = handle("");
SerialBlob blob = new SerialBlob(bytes);
System.out.println(name+product_name+"截图成功");
counts++;
System.err.println(name + "截取图片共计《" + counts + "》张");
String insert = "REPLACE INTO "+AttributeCommon.image_tablename+"(id,image,PROCESS_DATE) values(?,?,?)";
DatasHandle.insert(insert,id,blob,date,url);
} catch (Exception e) {
e.printStackTrace();
}
driver.get("http://www.baidu.com/");
try {
new File(path).delete();
} catch (Exception e) {
e.printStackTrace();
}
try {
new File(file_path).delete();
} catch (Exception e) {
e.printStackTrace();
}
}


/**
* 将byte数组写入图片文件中
* @param b
*/
private void bytesToImageFile(byte[] b){

try {
FileImageOutputStream imageOutput = new FileImageOutputStream(new File(file_path));
imageOutput.write(b, 0, b.length);
imageOutput.close();
} catch (Exception e) {
e.printStackTrace();
}
}

/**
* 重新加载压缩后的图片
* @return
* @throws IOException
*/
private byte[] handle(String path) throws IOException{

String real_path = "";
if(path==null || path.equals("")){
real_path = file_path;
}
else{
real_path = path;
}
BufferedInputStream in = null;
ByteArrayOutputStream out = null;
try {
in = new BufferedInputStream(new FileInputStream(new File(real_path)));
out = new ByteArrayOutputStream(1024);

byte[] temp = new byte[1024];
int size = 0;
while ((size = in.read(temp)) != -1) {
out.write(temp, 0, size);
}
try {
in.close();
} catch (Exception e) {
e.printStackTrace();
}
return out.toByteArray();
} catch (Exception e) {
e.printStackTrace();
return null;
} finally {
try {
out.close();
} catch (Exception e) {
e.printStackTrace();
}
}
}

/**
* Read an input stream into a string
* @param in
* @return
* @throws IOException
*/
public String streamToString(InputStream in) throws IOException {
StringBuffer out = new StringBuffer();
byte[] b = new byte[4096];
for (int n; (n = in.read(b)) != -1;) {
out.append(new String(b, 0, n));
}
return out.toString();
}


public byte[] takeScreenshot() throws IOException {
TakesScreenshot takesScreenshot = (TakesScreenshot) driver;
return takesScreenshot.getScreenshotAs(OutputType.BYTES);
}


public File takeScreenshot1() throws IOException {
TakesScreenshot takesScreenshot = (TakesScreenshot) driver;
return takesScreenshot.getScreenshotAs(OutputType.FILE);
}


public BufferedImage createElementImage(WebElement webElement)
throws IOException {
// 获得webElement的位置和大小。
Point location = webElement.getLocation();
Dimension size = webElement.getSize();
// 创建全屏截图。
BufferedImage originalImage = ImageIO.read(new ByteArrayInputStream(
takeScreenshot()));
// 截取webElement所在位置的子图。
BufferedImage croppedImage = originalImage.getSubimage(location.getX(),
location.getY(), size.getWidth(), size.getHeight());
return croppedImage;
}


public void snapshot(TakesScreenshot drivername, String filepath) {
// this method will take screen shot ,require two parameters ,one is
// driver name, another is file name


File scrFile = drivername.getScreenshotAs(OutputType.FILE);
// Now you can do whatever you need to do with it, for example copy
// somewhere
try {
System.out.println("save snapshot path is:" + filepath);
FileUtils.copyFile(scrFile, new File(filepath));
} catch (IOException e) {
System.out.println("Can't save screenshot");
e.printStackTrace();
} finally {
System.out.println("screen shot finished");
}
}

/**
* 初始化浏览器
*/
private void init() {


while (true) {
try {
System.out.println(name+"初始化浏览器");
System.setProperty("webdriver.firefox.bin",
PropertyUtil.getValue("path"));
driver = new FirefoxDriver();
driver.manage().timeouts().implicitlyWait(10, TimeUnit.SECONDS);
windowHandle = driver.getWindowHandle();
return;
} catch (Exception e) {
try {
driver.close();
} catch (Exception e2) {
e2.printStackTrace();
}
driver = null;
try {
Thread.sleep(1000);
} catch (InterruptedException e1) {
e1.printStackTrace();
}
}
}
}
}