AI 安全之

from requests_html import HTMLSession
import time

def save_links_to_file(links):
    # 追加链接到文件
    with open('urls.txt', 'a') as f:
        for link in links:
            f.write(link + '\n')

def crawl(url, end_time):
    # 创建HTMLSession对象，获取网页内容
    session = HTMLSession()
    r = session.get(url)
    # 获取所有链接
    links = r.html.absolute_links
    urls = []
    for link in links:
        # 如果链接是以http或https开头，则保存链接
        if link.startswith('http'):
            urls.append(link)
    # 保存链接到文件
    save_links_to_file(urls)
    # 递归访问所有链接
    for url in urls:
        # 如果当前时间已经超过了指定时间，则停止递归
        if time.time() > end_time:
            return
        crawl(url, end_time)


# 设置爬取的起始URL和结束时间
start_url = 'http://127.0.0.1:8000/'
end_time = time.time() + 3

# 调用函数，从指定URL开始递归访问所有链接，并追加到文件
crawl(start_url, end_time)

import os
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
import urllib.parse
def load(name):
    filepath = os.path.join(str(os.getcwd()), name)
    with open(filepath,'r') as f:
        alldata = f.readlines()
    ans = []
    for i in alldata:
        i = str(urllib.parse.unquote(i))
        ans.append(i)
    return ans

badqueries = load('badqueries.txt')#dao ru数据集
goodqueries = load('goodqueries.txt')#导入两类url
Y = [1 for i in range(0, len(badQueries))]+[0 for i in range(0, len(validQueries))]#给标签
vectorizer = TfidfVectorizer()#用来将url向量化
X = vectorizer.fit_transform(badqueries+goodqueries)#直接输进去
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42)#分割训练数据和测试数据
lgs = LogisticRegression(class_weight='balanced') #简单的逻辑回归二分类
lgs.fit(X_train, y_train)#begin！
print(lgs.score(X_test, y_test))
#0.9941273293313274

import re
from flask import Flask, request
import os
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
import urllib.parse
def load(name):
    filepath = os.path.join(str(os.getcwd()), name)
    with open(filepath,'r') as f:
        alldata = f.readlines()
    ans = []
    for i in alldata:
        i = str(urllib.parse.unquote(i))
        ans.append(i)
    return ans

badqueries = load('badqueries.txt')
goodqueries = load('goodqueries.txt')#导入两类url
Y = [1 for i in range(0, len(badqueries))]+[0 for i in range(0, len(goodqueries))]#给标签
vectorizer = TfidfVectorizer()#用来将url向量化
X = vectorizer.fit_transform(badqueries+goodqueries)#直接输进去
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42)#分割训练数据和测试数据
lgs = LogisticRegression(class_weight='balanced') #简单的逻辑回归二分类
lgs.fit(X_train, Y_train)#begin！
print(lgs.score(X_test, Y_test))
app = Flask(__name__)

@app.route('/check_url', methods=['POST'])
def check_url(): 
    url = request.form['url']
    X_predict = vectorizer.transform([url])
    res = lgs.predict(X_predict)
    if res==1:return '1\n'
    return '0\n'

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8889)

package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
)

func main() {
	// 启动HTTP服务器，监听8000端口
	http.HandleFunc("/", handleRequest)
	http.ListenAndServe(":8080", nil)
	println(1)
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	// 从请求中获取URL
	url := r.URL.Path[1:]
	// 发送HTTP请求到Python服务器
	result, err := checkURL(url)
	if err != nil {
		// 如果调用失败，返回500 Internal Server Error错误
		// 这块也可以优化，python服务器挂了等于go服务器也挂了，hhhh
		http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
		return
	}
	// 如果返回值为1，禁止访问该URL，否则允许访问
	if result == 1 {
		http.Error(w, "403 Forbidden", http.StatusForbidden)
	} else {
		fmt.Fprintf(w, "Hello, %s!", url)
	}
}

func checkURL(url string) (int, error) {
	// 构造HTTP请求
	data := fmt.Sprintf("url=%s", url)
	req, err := http.NewRequest("POST", "http://localhost:8889/check_url", strings.NewReader(data))
	if err != nil {
		return 0, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	// 发送HTTP请求
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	// 解析HTTP响应
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0, err
	}
	result := string(body)
	if result == "1\n" {
		return 1, nil
	} else {
		return 0, nil
	}
}