/***********************************************************************
Copyright Buu Information Security Major

File Name:secretsender

Author:姚小龙

ID:2021240381013

Version:3.0

Date:2023.7.1

Description:通过socket实现的tcp加密发送
************************************************************************/
#define _WIN32_WINNT  0x501 //NO causes error: undefined reference to `getaddrinfo'

#undef UNICODE//#undef指令用于“取消”已定义的#deinfe指令

#define WIN32_LEAN_AND_MEAN//指示编译器不要包含与MFC相关的操作

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <windows.h>   // WinApi header
#include <string.h>
#include <winsock2.h>
#include <ws2tcpip.h>

// Need to link with Ws2_32.lib, Mswsock.lib, and Advapi32.lib
#pragma comment (lib, "Ws2_32.lib")
#pragma comment (lib, "Mswsock.lib")
#pragma comment (lib, "AdvApi32.lib")

#define true 1
#define false 0
#define MsgVersion 3
#define DEFAULT_BUFLEN 512
#define DEFAULT_PORT "27013"

enum Type_msg
{
    SENDPAR=1,
    PAROK,
    My_number,
    Bye,
};

typedef struct _DH_PROTOCOL
{
    unsigned char version;
    enum Type_msg type;   //消息类型
    unsigned short p; //参数p
    unsigned short a; //参数a
    unsigned short YAB; //共享给对方的数值
    char Data[];
}DH_Protocol;

int server();//服务端
int Crypto(char*,char*);//明文
int Decode(char*,char*);//密文

int main(int argc, char *argv[])
{
        int IP_ADDRESS = argv[2];
        int Data = argv[4];

        //secretsender [options1(client or sever)] [ip-address] [options2] [Data]
        if(strcmp(argv[1], "-s") == NULL)//sever
                server();

        else if(strcmp(argv[1], "-c") == NULL)//client
        {
            if(strcmp(argv[3], "-n") == NULL)
                Crypto(IP_ADDRESS, Data);//明文直接显示

            if(strcmp(argv[3], "-m") == NULL)
                Decode(IP_ADDRESS, Data);//密文解密显示
        }

        else
        {
            printf("error\n");
            exit(1);
        }

        return 0;
}

/******************************************************************

Function Name: server

Parameters: 无

Return: 0

Description: 此函数是服务器函数，是服务器端口

*******************************************************************/
int server()
{
    WSADATA wsaData;
    int iResult;

    SOCKET ListenSocket = INVALID_SOCKET;
    SOCKET ClientSocket = INVALID_SOCKET;

    struct addrinfo *result = NULL,
                    hints;
    int iSendResult;
    //char receievebuffer[DEFAULT_BUFLEN]={'\0'};
    //int receievebufferlen = DEFAULT_BUFLEN;
    DH_Protocol *receievebuffer = (DH_Protocol *)malloc(DEFAULT_BUFLEN);
    DH_Protocol *sendbuffer = (DH_Protocol *)malloc(DEFAULT_BUFLEN);
    char temp[DEFAULT_BUFLEN];
    int receievebufferlen = DEFAULT_BUFLEN;
    // Initialize Winsock初始化Winsock
    iResult = WSAStartup(MAKEWORD(2,2), &wsaData);
    if (iResult != 0)
    {
        printf("WSAStartup failed with error: %d\n", iResult);
        return 1;
    }

    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;
    hints.ai_flags = AI_PASSIVE;
    // alert message from client
    printf("Waiting for connect...\n");

    // Resolve the server address and port
    iResult = getaddrinfo(NULL, DEFAULT_PORT, &hints, &result);
    if (iResult != 0)
    {
        printf("getaddrinfo failed with error: %d\n", iResult);
        WSACleanup();
        return 1;
    }
    // Create a SOCKET for connecting to server
    ListenSocket = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
    if (ListenSocket == INVALID_SOCKET)
    {
        printf("socket failed with error: %ld\n", (long)WSAGetLastError());
        freeaddrinfo(result);
        WSACleanup();
        return 1;
    }
    // Setup the TCP listening socket
    iResult = bind( ListenSocket, result->ai_addr, (int)result->ai_addrlen);
    if (iResult == SOCKET_ERROR)
    {
        printf("bind failed with error: %d\n", WSAGetLastError());
        freeaddrinfo(result);
        closesocket(ListenSocket);
        WSACleanup();
        return 1;
    }

    freeaddrinfo(result);

    iResult = listen(ListenSocket, SOMAXCONN);
    if (iResult == SOCKET_ERROR)
    {
        printf("listen failed with error: %d\n", WSAGetLastError());
        closesocket(ListenSocket);
        WSACleanup();
        return 1;
    }
    // Accept a client socket
    ClientSocket = accept(ListenSocket, NULL, NULL);

    if (ClientSocket == INVALID_SOCKET)
    {
        printf("accept failed with error: %d\n", WSAGetLastError());
        closesocket(ListenSocket);
        WSACleanup();
        return 1;
    }
    else
        printf("ok\n");

    // No longer need server socket
    closesocket(ListenSocket);
    // Receive until the peer shuts down the connection
    iResult = recv(ClientSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer, temp, DEFAULT_BUFLEN);
    if (iResult > 0)
    {
        sendbuffer->p = receievebuffer->p;
        sendbuffer->a = receievebuffer->a;
        sendbuffer->type = PAROK;
        iSendResult = send(ClientSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
        if (iSendResult == SOCKET_ERROR)
        {
            printf("send failed with error: %d\n", WSAGetLastError());
            closesocket(ClientSocket);
            WSACleanup();
            return 1;
        }
    }
    else
    {
        printf("recv failed with error: %d\n", WSAGetLastError());
        closesocket(ClientSocket);
        WSACleanup();
        return 1;
    }

    unsigned short key;
    sendbuffer->type = My_number;
    iResult = recv(ClientSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer, temp, DEFAULT_BUFLEN);
    /*if(receievebuffer->YAB != 0)
    {
        printf("enter your secret key:");
        scanf("%i",&key);
        sendbuffer->YAB = (int)(pow(sendbuffer->a, key) + 0.5) % sendbuffer->p;
    }
    else
        sendbuffer->YAB = 0;*/
    if (iResult > 0)
    {
        //发送YB
        iSendResult = send(ClientSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
        if(iSendResult == SOCKET_ERROR)
        {
            printf("send failed with error: %d\n", WSAGetLastError());
            closesocket(ClientSocket);
            WSACleanup();
            return 1;
        }
    }
    else
    {
        printf("recv failed with error: %d\n", WSAGetLastError());
        closesocket(ClientSocket);
        WSACleanup();
        return 1;
    }
    // Receive until the peer shuts down the connection
    iResult = recv(ClientSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer, temp, DEFAULT_BUFLEN);

    // shutdown the connection since we're done
    iResult = shutdown(ClientSocket, SD_SEND);
    if (iResult == SOCKET_ERROR)
    {
        printf("shutdown failed with error: %d\n", WSAGetLastError());
        closesocket(ClientSocket);
        WSACleanup();
        return 1;
    }


    int i;
    for(i = 0; i < (int)strlen(receievebuffer->Data); i++)
    {
        if(receievebuffer->Data[i] >= '0' && receievebuffer->Data[i] <= '9') receievebuffer->Data[i] = (receievebuffer->Data[i] - '0' + 10 - 3) % 10 + '0';
        else if(receievebuffer->Data[i] >= 'a' && receievebuffer->Data[i] <= 'z') receievebuffer->Data[i] = (receievebuffer->Data[i] - 'a' + 26 - 3) % 26 + 'a';
        else if(receievebuffer->Data[i] >= 'A' && receievebuffer->Data[i] <= 'Z') receievebuffer->Data[i] = (receievebuffer->Data[i] - 'A' + 26 - 3) % 26 + 'A';
    }
    receievebuffer->Data[i] = '\0';
    printf("received data :<%s>\n", receievebuffer->Data);
    // cleanup
    closesocket(ClientSocket);
    WSACleanup();
    free(receievebuffer);
    free(sendbuffer);
    return 0;
}


/******************************************************************

Function Name: Crypto

Parameters: char* IP_ADDRESS：连接用的IP地址
            char* Data：要传输的数据
Return: 0

Description: 此函数是明文传输函数

*******************************************************************/
int Crypto(char* IP_ADDRESS, char* Data)
{
   WSADATA wsaData;
    SOCKET ConnectSocket = INVALID_SOCKET;
    struct addrinfo *result = NULL, *ptr = NULL, hints;
    DH_Protocol *receievebuffer=(DH_Protocol *)malloc(DEFAULT_BUFLEN);
    DH_Protocol *sendbuffer=(DH_Protocol *)malloc(DEFAULT_BUFLEN);
    char temp[DEFAULT_BUFLEN];
    sendbuffer->version = MsgVersion;
    sendbuffer->type = SENDPAR;
    sendbuffer->p = 23;
    sendbuffer->a = 5;
    int iResult;
    int receievebufferlen = DEFAULT_BUFLEN;
	// Initialize winsock初始化Winsock
    iResult = WSAStartup(MAKEWORD(2,2), &wsaData);
    if (iResult != 0) {
        printf("WSAStartup failed with error: %d\n", iResult);
        return 1;
    }
    // Initialize Hints
    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;
    // Resolve the server address and port解析服务器地址和端口
    iResult = getaddrinfo(IP_ADDRESS, DEFAULT_PORT, &hints, &result);
    if ( iResult != 0 )
    {
        printf("getaddrinfo failed with error: %d\n", iResult);
        WSACleanup();
        return 1;
    }

    // Attempt to connect to an address until one succeeds尝试连接到一个地址，直到成功
    for(ptr = result; ptr != NULL ; ptr = ptr->ai_next)
    {
        // Create a SOCKET for connecting to server
        ConnectSocket = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (ConnectSocket == INVALID_SOCKET)
        {
            printf("socket failed with error: %ld\n", (long)WSAGetLastError());
            WSACleanup();
            return 1;
        }

        // Connect to server.连接到服务器
        iResult = connect(ConnectSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (iResult == SOCKET_ERROR)
        {
            closesocket(ConnectSocket);
            ConnectSocket = INVALID_SOCKET;
            continue;
        }
        break;
    }

    freeaddrinfo(result);

    if (ConnectSocket == INVALID_SOCKET)
    {
        printf("Unable to connect to server!\n");
        WSACleanup();
        return 1;
    }
    else
        printf("connect %s:ok\n", IP_ADDRESS);

    iResult = send(ConnectSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
    if (iResult == SOCKET_ERROR)
    {
        printf("send failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }
    iResult = recv(ConnectSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer, temp, DEFAULT_BUFLEN);//memcpy
    if (iResult > 0)
    {
        if(receievebuffer->type != PAROK)
        {
            printf("the peer receives public parameters failed\n");
            return 0;
        }
    }
    else
        printf("recv failed with error: %d\n", WSAGetLastError());

    //发送YA
    unsigned short key;

    sendbuffer->YAB = 0;
    sendbuffer->type = My_number;
    iResult = send(ConnectSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
    if (iResult == SOCKET_ERROR)
    {
        printf("send failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 0;
    }

    //接收YB
    iResult = recv(ConnectSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer ,temp, DEFAULT_BUFLEN);
    if (iResult > 0)
    {
        if(receievebuffer->type != My_number)
        {
            printf("the peer receives key failed\n");
            return 0;
        }
    }
    else
        printf("recv failed with error: %d\n", WSAGetLastError());

    //发送信息
    //unsigned short K = (int)(pow(receievebuffer->YAB, key) + 0.5) % sendbuffer->p;
    int i;
    for(i = 0; i < (int)strlen(Data); i++)
    {
        if(Data[i] >= '0' && Data[i] <= '9') sendbuffer->Data[i] = (Data[i] - '0' + 3) % 10 + '0';
        else if(Data[i] >= 'a' && Data[i] <= 'z') sendbuffer->Data[i] = (Data[i] - 'a' + 3) % 26 + 'a';
        else if(Data[i] >= 'A' && Data[i] <= 'Z') sendbuffer->Data[i] = (Data[i] - 'A' + 3) % 26 + 'A';
        else sendbuffer->Data[i] = Data[i];
    }
    sendbuffer->Data[i] = '\0';
    sendbuffer->type = Bye;
    printf("Data length:%d\n", (int)strlen(Data));
    printf("send Data <%s>:", Data);
    iResult = send(ConnectSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
    if (iResult == SOCKET_ERROR)
    {
        printf("send failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }
    printf("ok\n");

    // shutdown the connection since no more Data will be sent
    iResult = shutdown(ConnectSocket, SD_SEND);
    if (iResult == SOCKET_ERROR)
    {
        printf("shutdown failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }

    // Receive until the peer closes the connection

        if (iResult == 0)
            printf("Connection closed\n");
        else if(iResult < 0)
            printf("recv failed with error: %d\n", WSAGetLastError());


    // cleanup
    closesocket(ConnectSocket);
    WSACleanup();
    free(receievebuffer);
    free(sendbuffer);

    return 0;
}

/******************************************************************

Function Name: Decode

Parameters: char* IP_ADDRESS：连接用的IP地址
            char* Data：要传输的数据

Return: 0

Description: 此函数是密文传输函数

*******************************************************************/
int Decode(char* IP_ADDRESS, char* Data)
{
    WSADATA wsaData;
    SOCKET ConnectSocket = INVALID_SOCKET;

    struct addrinfo *result = NULL, *ptr = NULL, hints;
    DH_Protocol *receievebuffer = (DH_Protocol *)malloc(DEFAULT_BUFLEN);
    DH_Protocol *sendbuffer = (DH_Protocol *)malloc(DEFAULT_BUFLEN);
    char temp[DEFAULT_BUFLEN];
    sendbuffer->version = MsgVersion;
    sendbuffer->type = SENDPAR;
    sendbuffer->p = 23;
    sendbuffer->a = 5;
    int iResult;
    int receievebufferlen = DEFAULT_BUFLEN;

    // Initialize Winsock
    iResult = WSAStartup(MAKEWORD(2,2), &wsaData);//wsastartup
    if (iResult != 0)
    {
        printf("WSAStartup failed with error: %d\n", iResult);
        return 1;
    }
	ZeroMemory( &hints, sizeof(hints) );
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;
    // Resolve the server address and port
    iResult = getaddrinfo(IP_ADDRESS, DEFAULT_PORT, &hints, &result);
    if (iResult != 0)
    {
        printf("getaddrinfo failed with error: %d\n", iResult);
        WSACleanup();
        return 1;
    }
    // Attempt to connect to an address until one succeeds
    printf("connect %s:",IP_ADDRESS);
    for(ptr = result; ptr != NULL ;ptr = ptr->ai_next)
    {
        // Create a SOCKET for connecting to server
        ConnectSocket = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (ConnectSocket == INVALID_SOCKET)
        {
            printf("socket failed with error: %ld\n", (long)WSAGetLastError());
            WSACleanup();
            return 1;
        }
        // Connect to server.
        iResult = connect( ConnectSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (iResult == SOCKET_ERROR)
        {
            closesocket(ConnectSocket);
            ConnectSocket = INVALID_SOCKET;
            continue;
        }
        break;
    }
    freeaddrinfo(result);
    if (ConnectSocket == INVALID_SOCKET)
    {
        printf("Unable to connect to server!\n");
        WSACleanup();
        return 1;
    }
    printf("ok\n");
    iResult = send(ConnectSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
    if (iResult == SOCKET_ERROR)
    {
        printf("send failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }
    iResult = recv(ConnectSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer, temp, DEFAULT_BUFLEN);
    if (iResult > 0)
    {
        if(receievebuffer->type != PAROK)
        {
            printf("the peer receives public parameters failed\n");
            return 1;
        }
    }
    else
        printf("recv failed with error: %d\n", WSAGetLastError());
    //发送YA


    sendbuffer->type = My_number;
    iResult = send(ConnectSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
    if (iResult == SOCKET_ERROR)
    {
        printf("send failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }
    //接收YB
    iResult = recv(ConnectSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer, temp, DEFAULT_BUFLEN);
    if (iResult > 0)
    {
        if(receievebuffer->type != My_number)
        {
            printf("the peer receives key failed\n");
            return 1;
        }
    }
    else
        printf("recv failed with error: %d\n", WSAGetLastError());
    //发送信息
    int i;
    for(i = 0;i < (int)strlen(Data); i++)
    {
        if(Data[i] >= '0' && Data[i] <= '9') sendbuffer->Data[i] = (Data[i] - '0' + 3) % 10 + '0';
        else if(Data[i] >= 'a' && Data[i] <= 'z') sendbuffer->Data[i] = (Data[i] - 'a' + 3) % 26 + 'a';
        else if(Data[i] >= 'A' && Data[i] <= 'Z') sendbuffer->Data[i] = (Data[i] - 'A' + 3) % 26 + 'A';
        else sendbuffer->Data[i] = Data[i];
    }


    sendbuffer->Data[i] = '\0';
    sendbuffer->type = Bye;
    int k = i;
    printf("Data length:%d\n", (int)strlen(Data));
    printf("Plain <%s>\n", Data);
    printf("Cipher(hex)<");
    for(int i = 0; i < k - 1; i ++ ) printf("%x ", sendbuffer->Data[i]);
    printf("%x", sendbuffer->Data[k - 1]);
    printf(">\n");
    printf("send data Hex:<");
    for(int i = 0; i < k - 1; i ++ ) printf("%x ", sendbuffer->Data[i]);
    printf("%x", sendbuffer->Data[k - 1]);
    printf("> : ok\n");

    iResult = send(ConnectSocket, (char *)sendbuffer, DEFAULT_BUFLEN, 0);
    if (iResult == SOCKET_ERROR)
    {
        printf("send failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }
    // shutdown the connection since no more Data will be sent
    iResult = shutdown(ConnectSocket, SD_SEND);
    if (iResult == SOCKET_ERROR)
    {
        printf("shutdown failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }
    // Receive until the peer closes the connection

    iResult = recv(ConnectSocket, temp, receievebufferlen, 0);
    memcpy(receievebuffer, temp, DEFAULT_BUFLEN);

    if (iResult == 0)
        printf("Connection closed\n");
    else if(iResult < 0)
        printf("recv failed with error: %d\n", WSAGetLastError());

    // cleanup
    closesocket(ConnectSocket);
    WSACleanup();
    free(receievebuffer);
    free(sendbuffer);
    return 0;
}