WebsiteDetector::WebsiteDetector() 
{
	AddValidURL("http://jwc.hit.edu.cn/","http://jwts.hit.edu.cn/");
	AddBlockedURL("http://xltj.hit.edu.cn/");
}

class Cache {

public:

std::string GetDate(const std::string& url); // 獲取url對(duì)應(yīng)保存的LastModified字段

bool Get(const std::string& url, char* response, size_t& start, size_t& responseSize); // 讀取緩存

bool Put(const std::string& url, const char* response, size_t responseSize, size_t& start); // 保存緩存

private:

?      std::string cacheDirectory_; // 存放緩存的文件目錄

?      std::map<std::string, std::string> cacheMap_; // 對(duì)象和LastModified字段的映射關(guān)系

?      std::map<std::string, std::string> fileMap_; / 對(duì)象與文件名的映射關(guān)系

?      std::mutex mutex_; // 多線(xiàn)程同時(shí)讀寫(xiě)文件的互斥鎖

};

class HttpProxyServer

{

public:

	HttpProxyServer(int port); // 構(gòu)造函數(shù)，參數(shù)為端口號(hào)

	void Start(); // 監(jiān)聽(tīng)客戶(hù)端連接請(qǐng)求

private:

	int serverSocket_; // 代理服務(wù)Socket

	int port_; // 端口號(hào)

	struct sockaddr_in serverAddr_; // 代理服務(wù)地址

	Cache cache_; // Cache類(lèi)

	WebsiteDetector websiteDetector_; // websiteDetector類(lèi)

	void HandleClient(int clientSocket); // 子線(xiàn)程調(diào)用函數(shù)

	std::string ExtractUrl(const std::string &httpRequest); // 解析URL

	int CreateServerSocket(const std::string &host); // 創(chuàng)建與原服務(wù)器連接的Socket

	bool ServerToClient(const std::string &url, int clientSocket); // 轉(zhuǎn)發(fā)數(shù)據(jù)

	void ParseUrl(const std::string &url, std::string &host, std::string &path); // 解析主機(jī)名與路徑名

};

std::string locationResponse = std::string("HTTP/1.1 302 Found") + MY_CRLF + "Location: " + newUrl + MY_CRLF + MY_CRLF;

send(clientSocket, locationResponse.c_str(), locationResponse.size(), 0);

serverAddr_.sin_addr.s_addr = INADDR_ANY;

// serverAddr_.sin_addr.S_un.S_addr = inet_addr("127.0.0.1"); //只允許本機(jī)用戶(hù)訪(fǎng)問(wèn)服務(wù)器

date = cache_.GetDate(url);
std::string cacheRequest = httpRequest + "If-Modified-Since: " + date + MY_CRLF + MY_CRLF;

if (IsResponseNotModified(responseNotModified) ) {
?      // std::cout << "304 Not Modified" << std::endl;
?      sel = false;
}else {
?      cache_.ClearFileContent(url); //清空
?      sel = true;
}

//g++ your_code.cpp -o your_executable -lws2_32

#include <fstream>
#include <iostream>
#include <cstring>
#include <cstdlib>
#include <cstdio>
#include <string>
#include <thread>
#include <vector>
#include <mutex>
#include <list>
#include <map>
#include <sstream>
#include <winsock2.h>
#include <WS2tcpip.h>



#define MAX_CLIENTS 6
#define BUFSIZE 655360
#define HEADSIZE 128
#define MY_CRLF "\r\n"

class WebsiteDetector {
public:
    WebsiteDetector() 
    {
        AddValidURL("http://jwc.hit.edu.cn/", "http://jwts.hit.edu.cn/");
        AddBlockedURL("http://xltj.hit.edu.cn/");
    }
    // 釣魚(yú)
    std::string IsURLPhishing(const std::string& url) {
        auto it = validURLs_.find(url);
        if (it != validURLs_.end()) {
            return it->second;
        } else {
            return "Phishing";
        }
    }
    // 屏蔽
    bool IsURLBlocked(const std::string& url) {
        for (const std::string& blockedURL : blockedURLs_) {
             if (url.find(blockedURL) != std::string::npos) {
                return true;
            }
        }
        return false;
    }
private:
    std::map<std::string, std::string> validURLs_;
    std::vector<std::string> blockedURLs_;

    void AddValidURL(const std::string& srcURL, const std::string& dstURL) 
    {
        validURLs_[srcURL] = dstURL;
    }

    void AddBlockedURL(const std::string& url) {
        blockedURLs_.push_back(url);
    }
};

class Cache {
public:
   
    Cache() : cacheDirectory_("H:\\cppwork\\CS-networking\\.cache") {
        std::cout << cacheDirectory_ << std::endl;
        std::system(("mkdir -p " + cacheDirectory_).c_str());
    }
    bool Check(const std::string& url) {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = cacheMap_.find(url);
        if (it != cacheMap_.end()) 
        {
            return true;
        }
        return false;
    }
    // 清空文件內(nèi)容
    bool ClearFileContent(const std::string& url) {
        std::lock_guard<std::mutex> lock(mutex_);

        // Generate a unique filename based on the URL
        std::string fileName = GetFileNameFromUrl(url);
        auto it = fileMap_.find(fileName);
        std::string fileTag = it->second;

        std::string filePath = cacheDirectory_ + "\\" + fileTag;
        // 打開(kāi)文件并使用 std::ios::trunc 模式來(lái)清空文件內(nèi)容
        std::ofstream file(filePath, std::ios::out | std::ios::trunc);
        if (!file) {
            std::cerr << "無(wú)法打開(kāi)文件：" << filePath << std::endl;
            return false;
        }
        // 關(guān)閉文件
        file.close();
        return true;
    }
    std::string GetDate(const std::string& url) 
    {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = cacheMap_.find(url);
        return it->second;
    }
    bool Get(const std::string& url, char* response, size_t& start, size_t& responseSize) {
        std::lock_guard<std::mutex> lock(mutex_);

        // Generate a unique filename based on the URL
        std::string fileName = GetFileNameFromUrl(url);
        std::string fileTag = fileMap_[fileName];

        std::cout << "Get() url: " << url << std::endl;
        std::cout << "Get() fileTag: " << fileTag << std::endl;
        // If found, read the response from the file
        std::ifstream file(cacheDirectory_ + "\\" + fileTag, std::ios::binary);
        if (file) {
            file.seekg(start, std::ios::beg);
            file.read(response, BUFSIZE);
            // Get the number of bytes read in this chunk
            size_t bytesRead = static_cast<size_t>(file.gcount());
            start += bytesRead;
            responseSize = bytesRead;
            response[bytesRead] = '\0';
            file.close();
            return true;
        }
        return false; // URL not found in the cache
    }

    bool Put(const std::string& url, const char* response, size_t responseSize, size_t& start) {
        std::lock_guard<std::mutex> lock(mutex_);

        // Generate a unique filename based on the URL
        std::string fileName = GetFileNameFromUrl(url);
        auto it = fileMap_.find(fileName);
        std::string fileTag;
        if (it == fileMap_.end()) 
        {
            fileTag = std::to_string(cnt);
            fileMap_[fileName] = fileTag;
            cnt ++;
        }else 
        {
            fileTag = it->second;
        }

        // Store the response in a file
        std::ofstream file(cacheDirectory_ + "\\" + fileTag, std::ios::binary | std::ios::app );
        if (!file) {
            fprintf(stderr, "file open error: %s(errno: %d)\n", strerror(errno),errno);
            return false; // Unable to open file for writing
        }
        file.seekp(start);
        file.write(response, responseSize);
        file.close();
        start += responseSize;
        return true; // Failed to store response in the cache
    }

    void PutDate(const std::string& url)
    {
        std::lock_guard<std::mutex> lock(mutex_);
        // Generate a unique filename based on the URL
        std::string fileName = GetFileNameFromUrl(url);
        std::string fileTag = fileMap_[fileName];
        // 拼接完整的文件路徑
        std::string filePath = cacheDirectory_ + "\\" + fileTag;

        // 打開(kāi)文件并讀取 Last-Modified 首部?jī)?nèi)容
        std::ifstream file(filePath);
        if (!file) {
            fprintf(stderr, "file open error: %s(errno: %d)\n", strerror(errno),errno);
        }
        std::string line;
        while (std::getline(file, line)) {
            // 查找包含 Last-Modified 首部的行
            if (line.find("Last-Modified:") != std::string::npos) {
                // 提取 Last-Modified 的值并存儲(chǔ)到 cacheMap_
                size_t startPos = line.find(":") + 2;
                size_t endPos = line.find(MY_CRLF);
                std::string date = line.substr(startPos, endPos);
                // std::cout << "line: " << line << std::endl;
                // std::cout << "date: " << date << std::endl;
                cacheMap_[url] = date;
                break; // 找到后可以退出循環(huán)
            }else
            {
                if (line == MY_CRLF)
                {   
                    break;
                }
            }
        }
        file.close();
    }

    
private:
    std::string cacheDirectory_;
    std::map<std::string, std::string> cacheMap_;
    std::map<std::string, std::string> fileMap_;
    std::mutex mutex_;
    int cnt = 1;

    std::string GetFileNameFromUrl(const std::string& url) {
        // Replace characters in the URL to create a valid filename
        std::string fileName = url;
        for (char& c : fileName) {
            if (c == '/' || c == '?' || c == '&' || c == '=') {
                c = '_';
            }
        }
        return fileName;
    }
};


// 定義HTTP代理服務(wù)器類(lèi)
class HttpProxyServer
{
public:
    HttpProxyServer(int port) : port_(port)
    {
        
        // 初始化服務(wù)器
        // 創(chuàng)建主套接字并綁定端口
        serverSocket_ = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket_ == -1)
        {
            fprintf(stderr, "Constructor(): create socket error: %s(errno: %d)\n", strerror(errno),errno);
            exit(EXIT_FAILURE);
        }

        // 設(shè)置服務(wù)器地址信息
        // 初始化 serverAddr_
        memset(&serverAddr_, 0, sizeof(serverAddr_));
        serverAddr_.sin_family = AF_INET;
        serverAddr_.sin_addr.s_addr = INADDR_ANY;
        // serverAddr_.sin_addr.S_un.S_addr = inet_addr("127.0.0.1"); //只允許本機(jī)用戶(hù)訪(fǎng)問(wèn)服務(wù)器
        serverAddr_.sin_port = htons(port_);

        // 綁定套接字到指定端口
        if (bind(serverSocket_, (struct sockaddr *)&serverAddr_, sizeof(serverAddr_)) == -1)
        {
            fprintf(stderr, "Constructor(): bind socket error: %s(errno: %d)\n",strerror(errno), errno);
            closesocket(serverSocket_);
            exit(EXIT_FAILURE);
        }

        // 開(kāi)始監(jiān)聽(tīng)客戶(hù)端連接請(qǐng)求
        if (listen(serverSocket_, MAX_CLIENTS) == -1)
        {
            fprintf(stderr, "Constructor(): listen socket error: %s(errno: %d)\n",strerror(errno),errno);
            closesocket(serverSocket_);
            exit(EXIT_FAILURE);
        }

        std::cout << "Proxy server started on port " << port_ << std::endl;
    }

    void Start()
    {
        // 啟動(dòng)服務(wù)器，監(jiān)聽(tīng)客戶(hù)端連接請(qǐng)求
        while (true)
        {
            struct sockaddr_in clientAddr;
            int clientAddrLen = sizeof(struct sockaddr);

            // 接受客戶(hù)端連接
            int clientSocket = accept(serverSocket_, (struct sockaddr *)&clientAddr, &clientAddrLen);
            if (clientSocket == INVALID_SOCKET)
            {
                fprintf(stderr, "Start(): accept socket error: %s(errno: %d)",strerror(errno),errno);
                continue; // 繼續(xù)等待下一個(gè)連接
            }

            // std::cout << "Start(): Accepted a client connection" << std::endl;

            // 創(chuàng)建子線(xiàn)程處理客戶(hù)端請(qǐng)求
            std::thread clientThread(&HttpProxyServer::HandleClient, this, clientSocket);
            clientThread.detach(); // 不等待
        }
    }

private:
    int serverSocket_;
    int port_;
    struct sockaddr_in serverAddr_;
    Cache cache_;
    WebsiteDetector websiteDetector_;

    void HandleClient(int clientSocket)
    {
        // 讀取客戶(hù)端的HTTP請(qǐng)求
        char buffer[BUFSIZE];
        memset(buffer, 0, BUFSIZE);
        ssize_t bytesRead = recv(clientSocket, buffer, BUFSIZE - 1, 0);
        if (bytesRead == -1)
        {
            perror("HandleClient(): Error reading from client socket");
            closesocket(clientSocket);
            return;
        }

        // 解析請(qǐng)求，提取URL
        std::string request(buffer);
        std::string url = ExtractUrl(request);

        std::cout << "<" << url << ">" << std::endl;
        
        // Website Filter; User Filter ; Website phishing
        if (websiteDetector_.IsURLBlocked(url))
        {
            std::cout << "Url Blocked Success: " << url << std::endl;
        }else
        {
            std::string newUrl = websiteDetector_.IsURLPhishing(url);
            if (newUrl == "Phishing")
            {
                // 向服務(wù)端請(qǐng)求，向客戶(hù)端發(fā)送
                if( ServerToClient(url, clientSocket) )
                {
                    std::cout << "Transmit Success!" << std::endl;
                }else
                {
                    std::cout << "Transmit Fail!" << std::endl;
                }
            }else
            {
                std::cout << "Phishing" << std::endl;
                std::string locationResponse = std::string("HTTP/1.1 302 Found") + MY_CRLF + "Location: " + newUrl + MY_CRLF + MY_CRLF;
                std::cout << locationResponse << std::endl;
                send(clientSocket, locationResponse.c_str(), locationResponse.size(), 0);
            }
            
        }

        std::cout << "----------------------" << std::endl;
        // 關(guān)閉連接
        closesocket(clientSocket);
    }

    // 提取URL
    std::string ExtractUrl(const std::string &httpRequest)
    {
        std::string url;
        // Debug
        // std::cout << "ExtractUrl(): httpRequest = " << std::endl << httpRequest << std::endl;

        // 在HTTP請(qǐng)求中查找"GET "，通常URL緊隨其后
        size_t getPos = httpRequest.find("GET ");
        if (getPos != std::string::npos)
        {
            // 找到"GET "后，查找下一個(gè)空格，該空格之后是URL
            size_t spacePos = httpRequest.find(' ', getPos + 4);
            if (spacePos != std::string::npos)
            {
                url = httpRequest.substr(getPos + 4, spacePos - (getPos + 4));
            }
        }
        return url;
    }
    
    void ParseUrl(const std::string &url, std::string &host, std::string &path)
    {
        // 查找 URL 中的 "http://"，并獲取其后的部分
        size_t httpPos = url.find("http://");
        if (httpPos != std::string::npos)
        {
            std::string urlWithoutHttp = url.substr(httpPos + 7); // 7 是 "http://" 的長(zhǎng)度
            // 查找 "/"，分隔主機(jī)名和路徑
            size_t slashPos = urlWithoutHttp.find('/');
            if (slashPos != std::string::npos)
            {
                host = urlWithoutHttp.substr(0, slashPos);
                path = urlWithoutHttp.substr(slashPos);
            }
            else
            {
                // 如果沒(méi)有找到 "/"，則整個(gè)剩余部分都是主機(jī)名
                host = urlWithoutHttp;
                path = "/";
            }
        }
        else
        {
            // 如果沒(méi)有 "http://" 前綴，則默認(rèn)協(xié)議為 HTTP，整個(gè) URL 都是主機(jī)名
            host = url;
            path = "/";
        }

        // Debug
        // std::cout << "url: " + url << std::endl;
        // std::cout << "host: " + host << std::endl;
        // std::cout << "path: " + path << std::endl;

    }
    int CreateServerSocket(const std::string &host)
    {
        // 域名解析
        addrinfo* result = NULL;
        addrinfo hints;

        ZeroMemory(&hints, sizeof(hints));
        hints.ai_family = AF_INET;  // 使用IPv4地址
        hints.ai_socktype = SOCK_STREAM;

        if (getaddrinfo(host.c_str(), "http", &hints, &result) != 0)
        {
            fprintf(stderr, "CreateServerSocket(): Failed to resolve the host: %s\n", host.c_str());
            return -1; // 返回-1表示連接失敗
        }

        // 創(chuàng)建Socket
        int serverSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket == -1)
        {
            fprintf(stderr, "CreateServerSocket(): create socket error: %s(errno: %d)\n", strerror(errno), errno);
            freeaddrinfo(result);  // 釋放內(nèi)存
            return -1; // 返回-1表示連接失敗
        }

        // 設(shè)置服務(wù)器地址信息
        struct sockaddr_in serverAddr;
        memset(&serverAddr, 0, sizeof(serverAddr));
        serverAddr.sin_family = AF_INET;
        serverAddr.sin_port = htons(80); // 設(shè)置端口號(hào)為80，可以根據(jù)需要修改
        serverAddr.sin_addr.s_addr = ((struct sockaddr_in *)(result->ai_addr))->sin_addr.s_addr;

        // 連接到原服務(wù)器
        if (connect(serverSocket, (struct sockaddr *)&serverAddr, sizeof(serverAddr)) == -1)
        {
            fprintf(stderr, "CreateServerSocket(): connect error: %s(errno: %d)\n",strerror(errno),errno);
            closesocket(serverSocket); // 在Windows中使用closesocket關(guān)閉套接字
            freeaddrinfo(result);  // 釋放內(nèi)存
            return -1; // 返回-1表示連接失敗
        }
        freeaddrinfo(result);  // 釋放內(nèi)存
        return serverSocket; // 返回連接成功的套接字描述符
    }

    bool ServerToClient(const std::string &url, int clientSocket)
    {
        // 解析URL，獲取主機(jī)名和路徑
        std::string host, path;
        ParseUrl(url, host, path);

        // 創(chuàng)建Socket連接到原服務(wù)器
        int serverSocket = CreateServerSocket(host);
        if (serverSocket == -1)
        {
            return FALSE; // 處理連接失敗的情況
        }  
        // 構(gòu)建HTTP請(qǐng)求
        std::string httpRequest = "GET " + path + " HTTP/1.1" + MY_CRLF + "Host: " + host + MY_CRLF + "Connection: close" + MY_CRLF;

        std::string date;
        bool sel;

        if (cache_.Check(url)) 
        {
            sel = false;
            date = cache_.GetDate(url);
            std::string cacheRequest = httpRequest + "If-Modified-Since: " + date + MY_CRLF + MY_CRLF;
            // 發(fā)送HTTP, 帶有If-Modified-Since 首部行
            if (send(serverSocket, cacheRequest.c_str(), cacheRequest.size(), 0) == -1)
            {
                perror("Error sending request to server");
                closesocket(serverSocket);
                return FALSE;
            }
            std::string cacheResponse;
            char cacheBuffer[HEADSIZE];
            ssize_t cacheBytesRead;
            cacheBytesRead = recv(serverSocket, cacheBuffer, HEADSIZE - 1, 0);
            std::string responseNotModified(cacheBuffer);

    
            // std::cout << "responseNotModified: " << responseNotModified << std::endl;

            if (IsResponseNotModified(responseNotModified) ) 
            {
                // std::cout << "304 Not Modified" << std::endl;
                sel = false;
            }else
            {
                cache_.ClearFileContent(url); //清空
                sel = true;
            }

        }else
        {
            sel = true;
        }
    
        if (sel == false) 
        {

            // std::cout << "cache hit!" << std::endl;
            // 接收緩存，轉(zhuǎn)發(fā)到客戶(hù)端
            char buffer[BUFSIZE];
            size_t start = 0;
            size_t bytesRead;
            while (1)
            {
                if (cache_.Get(url, buffer, start, bytesRead) == false)
                {
                    perror("Error sending response to client");
                }

                // std::cout << "bytesRead: " << bytesRead << std::endl;
                if (bytesRead == 0) break;

                if (send(clientSocket, buffer, bytesRead, 0) == -1)
                {
                    perror("Error sending response to client");
                    closesocket(serverSocket);
                    return FALSE;
                }
            }
           
        }else
        {
            httpRequest += MY_CRLF;
            // 發(fā)送HTTP請(qǐng)求到原服務(wù)器
            if (send(serverSocket, httpRequest.c_str(), httpRequest.size(), 0) == -1)
            {
                perror("Error sending request to server");
                closesocket(serverSocket);
                return FALSE;
            }

            // 接收原服務(wù)器的HTTP響應(yīng)
            char buffer[BUFSIZE];
            size_t start = 0;
            ssize_t bytesRead;
            while ((bytesRead = recv(serverSocket, buffer, BUFSIZE - 1, 0)) > 0)
            {
                buffer[bytesRead] = '\0';
                // 發(fā)送接收到的數(shù)據(jù)到客戶(hù)端
                if (send(clientSocket, buffer, bytesRead, 0) == -1)
                {
                    perror("Error sending response to client");
                    closesocket(serverSocket);
                    return FALSE;
                }
                if(cache_.Put(url, buffer, bytesRead, start) == false)
                {
                    std::cerr << "Cache put error" << std::endl;
                }
            }
            cache_.PutDate(url);
            if (! cache_.Check(url))
            {
                cache_.ClearFileContent(url);
            }
        }
        // 關(guān)閉原服務(wù)器連接
        closesocket(serverSocket);
        return TRUE;
    }
    bool IsResponseNotModified(const std::string& response) {
       // 查找第一個(gè)空格，定位到狀態(tài)碼的開(kāi)始
        size_t spacePos = response.find(' ');
        if (spacePos != std::string::npos) {
            // 提取狀態(tài)碼部分
            std::string statusCode = response.substr(spacePos + 1, 3);
            // 檢查狀態(tài)碼是否為 "304"
            return (statusCode == "304"); // HTTP/1.1 304 Not Modified
        }
        return false; // 未找到狀態(tài)碼
    }
};



bool InitWinsock()
{
    // 加載套接字庫(kù)（必須）
    WORD wVersionRequested;
    WSADATA wsaData;
    // 套接字加載時(shí)錯(cuò)誤提示
    int err;
    // 版本 2.2
    wVersionRequested = MAKEWORD(2, 2);
    // 加載 dll 文件 Scoket 庫(kù)
    err = WSAStartup(wVersionRequested, &wsaData);
    if (err != 0)
    {
        // 找不到 winsock.dll
        printf("加載 winsock 失敗，錯(cuò)誤代碼為: %d\n", WSAGetLastError());
        return FALSE;
    }
    if (LOBYTE(wsaData.wVersion) != 2 || HIBYTE(wsaData.wVersion) != 2)
    {
        printf("不能找到正確的 winsock 版本\n");
        return FALSE;
    }

    return TRUE;
}

int main()
{
    if (!InitWinsock())
    {
        WSACleanup();
        return -1; // 初始化失敗，退出程序
    }
    int port = 666; // 設(shè)置端口
    HttpProxyServer proxyServer(port);
    proxyServer.Start();
    WSACleanup(); // 在程序結(jié)束時(shí)清理Winsock庫(kù)
    return 0;
}