#include "StdAfx.h"
#include "WinSsh2Proc.h"

CWinSsh2Proc::CWinSsh2Proc(void)
{

}

CWinSsh2Proc::~CWinSsh2Proc(void)
{

}

bool CWinSsh2Proc::InitSocket()
{
	WSADATA wsadata;
	int err = WSAStartup(MAKEWORD(2, 0), &wsadata);
	if( err != 0 ) {
		fprintf(stderr, "WSAStartup failed with error: %d\n", err);
		return false;
	}

	return true;
}

int CWinSsh2Proc::waitsocket(int socket_fd, LIBSSH2_SESSION *session)
{
	struct timeval timeout;
	int rc;
	fd_set fd;
	fd_set *writefd = NULL;
	fd_set *readfd = NULL;
	int dir;

	timeout.tv_sec = 10;
	timeout.tv_usec = 0;

	FD_ZERO(&fd);

	FD_SET(socket_fd, &fd);

	/* now make sure we wait in the correct direction */
	dir = libssh2_session_block_directions(session);

	if(dir & LIBSSH2_SESSION_BLOCK_INBOUND)
		readfd = &fd;

	if(dir & LIBSSH2_SESSION_BLOCK_OUTBOUND)
		writefd = &fd;

	rc = select(socket_fd + 1, readfd, writefd, NULL, &timeout);

	return rc;
}

bool CWinSsh2Proc::ssh2_connect(std::string host, std::string user, std::string pwd, int port /*=22*/)
{
	m_user = user;
	m_pwd = pwd;
	m_port = port;
	m_host = host;
	// 创建socket连接;
	unsigned long hostaddr = inet_addr(host.c_str());
	SOCKET sock = socket(AF_INET, SOCK_STREAM, 0);
	m_sin.sin_family = AF_INET;
	m_sin.sin_port = htons(port);
	m_sin.sin_addr.s_addr = hostaddr;
	if ( connect(sock, (const sockaddr*)&m_sin, sizeof(m_sin)) != 0 ) {
		printf("failed to connect!\n");
		return false;
	}

	closesocket(sock);

	return true;
}

/*
void CWinSsh2Proc::ssh2_disconnect()
{
	int rc;
	int bytecount = 0;
	int exitcode = 127;
	char *exitsignal = (char *)"none";
	m_bStatus = false;
	if ( m_pChannel == NULL )
		goto shutdown;

	while((rc = libssh2_channel_close(m_pChannel)) == LIBSSH2_ERROR_EAGAIN)
		waitsocket(m_sock, m_pSession);

	if(rc == 0) {
		exitcode = libssh2_channel_get_exit_status(m_pChannel);
		libssh2_channel_get_exit_signal(m_pChannel, &exitsignal, NULL, NULL, NULL, NULL, NULL);
	}

	if(exitsignal)
		fprintf(stderr, "\nGot signal: %s\n", exitsignal);
	else
		fprintf(stderr, "\nEXIT: %d bytecount: %d\n", exitcode, bytecount);

	libssh2_channel_free(m_pChannel);
	m_pChannel = NULL;

shutdown:
	libssh2_session_disconnect(m_pSession, "Normal Shutdown, Thank you for playing");
	libssh2_session_free(m_pSession);
	m_pSession = NULL;

#ifdef WIN32
	closesocket(m_sock);
#else
	close(m_sock);
#endif
	fprintf(stderr, "all done\n");

	libssh2_exit();
}
*/

bool CWinSsh2Proc::ssh2_execute_command(std::string cmd, std::string &result, DWORD dwTimeout /*= 1000*/)
{
	int rc;
	bool bRet = false;
	SOCKET sock = INVALID_SOCKET;
	LIBSSH2_SESSION* pSession = NULL;
	LIBSSH2_CHANNEL* pChannel = NULL;
	if ( !ssh2_init_session(&pSession, sock) )
		return false;

	// 请求一个shell;
	while((pChannel= libssh2_channel_open_session(pSession)) == NULL && 
		libssh2_session_last_error(pSession, NULL, NULL, 0) == LIBSSH2_ERROR_EAGAIN) {
		waitsocket(sock, pSession);
	}

	if ( pChannel == NULL ) {
		printf("Error\n");
		goto shutdown;
	}
#if 0
	// 请求一个模拟终端;
	if( libssh2_channel_request_pty(pChannel,"VT100") != LIBSSH2_ERROR_NONE ) {
		printf("ssh请求伪终端失败，退出。\n");
		goto skip_shell;
	}

	// 在模拟终端上开户shell;
	if(libssh2_channel_shell(pChannel) != LIBSSH2_ERROR_NONE ) {
		printf("ssh shell请求失败，退出\n");
		goto shutdown;
	}
#else	
	while((rc = libssh2_channel_exec(pChannel, cmd.c_str())) == LIBSSH2_ERROR_EAGAIN) {
			waitsocket(sock, pSession);
	}
	if ( rc != LIBSSH2_ERROR_NONE )
		goto skip_shell;
#endif

	// 设置为阻塞模式;
	libssh2_channel_set_blocking(pChannel,0);

	if ( cmd.at(cmd.size()-1) != '\n' )
		cmd.append("\n");

	while((rc = libssh2_channel_write(pChannel, cmd.c_str(),  cmd.size())) == LIBSSH2_ERROR_EAGAIN ) {
		printf("ssh写堵塞，重试\n");
		waitsocket(sock, pSession);
	}

	if ( rc < 0 )
		goto shutdown;

	Sleep(dwTimeout);
	ssh2_read_channel(pChannel, result);
	bRet = true;

skip_shell:
	if(pChannel) {
		libssh2_channel_free(pChannel);
		pChannel = NULL;
	}

shutdown:
	libssh2_session_disconnect(pSession,"Normal Shutdown, Thank you for playing");
	libssh2_session_free(pSession);

#ifdef WIN32
	closesocket(sock);
#else
	close(sock);
#endif
	libssh2_exit();

	return bRet;
}

void CWinSsh2Proc::ssh2_read_channel(LIBSSH2_CHANNEL* pChannel, std::string &buffer)
{
	if ( pChannel == NULL )
		return;

	int rc;
	char buf[1024] = "";
	buffer.clear();
	do {
		rc = libssh2_channel_read(pChannel, buf, sizeof(buf) - 1);
		if(rc > 0) {
			buf[rc] = 0;
			buffer.append(buf, rc);
			memset(buf, 0, sizeof(buf));
		}
	}while( rc > 0 );

	printf("%s \n", buffer.c_str());
}

bool CWinSsh2Proc::ssh2_init_session(LIBSSH2_SESSION** pSession, SOCKET &sock)
{
	// 初始化库;
	int rc = libssh2_init(0);
	if ( rc != LIBSSH2_ERROR_NONE ) {
		fprintf(stderr, "libssh2 initialization failed (%d)\n", rc);
		return false;
	}

	sock = socket(AF_INET, SOCK_STREAM, 0);
	if ( connect(sock, (const sockaddr*)&m_sin, sizeof(m_sin)) != 0 ) {
		fprintf(stderr, "failed to connect!\n");
		return false;
	}

	// 建立ssh会话实例;
	*pSession = libssh2_session_init();
	libssh2_session_set_blocking(*pSession, 1);

	rc = libssh2_session_handshake(*pSession, sock);
	if ( rc != LIBSSH2_ERROR_NONE ){
		printf("Failure establishing SSH session: %d\n", rc);
		// 释放资源;
		SHUTDOWN(*pSession, sock)
		return false;
	}

	// 登录linux;
	rc = libssh2_userauth_password(*pSession, m_user.c_str(), m_pwd.c_str());
	if( rc != LIBSSH2_ERROR_NONE ) {
		printf("Authentication by password failed: %d\n", rc);
		// 释放资源;
		SHUTDOWN(*pSession, sock)
		return false;
	}
	
	return true;
}

bool CWinSsh2Proc::ssh2_sftp_download(std::string sftppath, std::string localpath)
{
	SOCKET sock = -1;
	LIBSSH2_SESSION *session;
	LIBSSH2_SFTP *sftp_session;
	LIBSSH2_SFTP_HANDLE *sftp_handle;
	// 初始化库;
	int rc = libssh2_init(0);
	if ( rc != LIBSSH2_ERROR_NONE ) {
		fprintf(stderr, "libssh2 initialization failed (%d)\n", rc);
		return false;
	}

	// 创建socket连接;
	sock = socket(AF_INET, SOCK_STREAM, 0);
	if ( connect(sock, (const sockaddr*)&m_sin, sizeof(m_sin)) != 0 ) {
		fprintf(stderr, "failed to connect!\n");
		return false;
	}

	// 建立ssh会话实例;
	session = libssh2_session_init();
	/* Since we have set non-blocking, tell libssh2 we are blocking */
	libssh2_session_set_blocking(session, 1);

	//while ( LIBSSH2_ERROR_EAGAIN == (rc = libssh2_session_handshake(m_pSession, m_sock)) );
	rc = libssh2_session_handshake(session, sock);
	if ( rc != LIBSSH2_ERROR_NONE ){
		fprintf(stderr, "Failure establishing SSH session: %d\n", rc);
		return false;
	}

	// 登录linux;
	bool bDone = false;
	//while ( LIBSSH2_ERROR_EAGAIN == (rc = libssh2_userauth_password(m_pSession, user.c_str(), pwd.c_str())));
	rc = libssh2_userauth_password(session, m_user.c_str(), m_pwd.c_str());
	if( rc != LIBSSH2_ERROR_NONE ) {
		fprintf(stderr, "Authentication by password failed: %d\n", rc);
		goto shutdown;
	}

	sftp_session = libssh2_sftp_init(session);
	if (!sftp_session) {
		goto shutdown;
	}

	/* Request a file via SFTP */
	sftp_handle = libssh2_sftp_open(sftp_session, sftppath.c_str(), LIBSSH2_FXF_READ, 0);


	if (!sftp_handle) {
		goto shutdown;
	}

	FILE *stream;
	if (fopen_s(&stream, localpath.c_str(), "wb") == 0) {
		do {
			char mem[1024];

			/* loop until we fail */
			rc = libssh2_sftp_read(sftp_handle, mem, sizeof(mem));

			if (rc > 0) {
				//从内存到磁盘
				fwrite(mem, 1, rc, stream);
			}
			else {
				break;
			}
		} while (1);

		bDone = true;
		fclose(stream);
	}


	libssh2_sftp_close(sftp_handle);
	libssh2_sftp_shutdown(sftp_session);

shutdown:
	libssh2_session_disconnect(session, "Normal Shutdown, Thank you for playing");
	libssh2_session_free(session);
	closesocket(sock);//INVALID_SOCKET
	
	return bDone;
}