package main

import (
	"fmt"
	"io"
	"net"
	"os"
)

// Proxy struct to hold connection data
type Proxy struct {
	Conn *net.TCPConn
	Addr *net.TCPAddr
	err  error
}

func (p *Proxy) SetAddrAndPort(ip string, port int) {
	address := fmt.Sprintf("%s:%d", ip, port)
	p.Addr, p.err = net.ResolveTCPAddr("tcp4", address)
	p.CheckError()

}

// OpenConnection opens dials the tcp connection
func (p *Proxy) OpenConnection() {
	p.Conn, p.err = net.DialTCP("tcp4", nil, p.Addr)
	p.CheckError()

}

func (p *Proxy) WriteBytes(ch chan []byte) {

	select {
	case d := <-ch:
		p.Conn.Write(d)
	}
	p.CheckError()

}

func (p *Proxy) ReadAll(ch chan []byte) {

	data := make([]byte, 0, 16384)
	tmp := make([]byte, 256)
	for {
		size, err := p.Conn.Read(tmp)

		if err != nil {
			if err == io.EOF {
				break
			}
			fmt.Printf("\n\n\n### Error on read %v", err)
			os.Exit(1)
		}
		data = append(data, tmp[:size]...)
		if size < 256 {
			ch <- data
		}
	}
	p.CheckError()

}

func (p *Proxy) CheckError() {
	if p.err != nil {
		fmt.Println(p.err)
		os.Exit(1)
	}
}

func NewProxy() *Proxy {
	return &Proxy{}
}

func ReplaceNonPrintable(data []byte) []byte {
	for i, b := range data {
		if (b < 32 && b != 13 && b != 10) || b > 127 {
			data[i] = 0
		}
	}
	return data
}