diff --git a/.idea/dataSources.xml b/.idea/dataSources.xml new file mode 100644 index 0000000..8f50f6f --- /dev/null +++ b/.idea/dataSources.xml @@ -0,0 +1,17 @@ + + + + + postgresql + true + org.postgresql.Driver + jdbc:postgresql://ep-jolly-king-a1xh98wj-pooler.ap-southeast-1.aws.neon.tech:5432/neondb + + + + + + $ProjectFileDir$ + + + \ No newline at end of file diff --git a/db/main.go b/db/main.go new file mode 100644 index 0000000..544ab4b --- /dev/null +++ b/db/main.go @@ -0,0 +1,35 @@ +package db + +import ( + "context" + "fmt" + "time" + + "database/sql" + + "github.com/uptrace/bun" + "github.com/uptrace/bun/dialect/pgdialect" + "github.com/uptrace/bun/driver/pgdriver" +) + +func GetDB(pgDsn string) (*bun.DB, error) { + if pgDsn == "" { + return nil, fmt.Errorf("pg_dsn is required") + } + + sqlConnection := sql.OpenDB(pgdriver.NewConnector( + pgdriver.WithDSN(pgDsn), + )) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := sqlConnection.PingContext(ctx); err != nil { + _ = sqlConnection.Close() + return nil, fmt.Errorf("failed to connect to database: %w", err) + } + + db := bun.NewDB(sqlConnection, pgdialect.New()) + + return db, nil +} diff --git a/go.mod b/go.mod index 059311d..b21e145 100644 --- a/go.mod +++ b/go.mod @@ -2,13 +2,26 @@ module sshchat go 1.25 -require github.com/gliderlabs/ssh v0.3.8 +require ( + github.com/gliderlabs/ssh v0.3.8 + github.com/joho/godotenv v1.5.1 + github.com/oschwald/geoip2-golang v1.13.0 + github.com/uptrace/bun v1.2.15 + github.com/uptrace/bun/dialect/pgdialect v1.2.15 + github.com/uptrace/bun/driver/pgdriver v1.2.15 + golang.org/x/crypto v0.43.0 +) require ( github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect - github.com/joho/godotenv v1.5.1 // indirect - github.com/oschwald/geoip2-golang v1.13.0 // indirect + github.com/jinzhu/inflection v1.0.0 // indirect github.com/oschwald/maxminddb-golang v1.13.0 // indirect - golang.org/x/crypto v0.43.0 // indirect + github.com/puzpuzpuz/xsync/v3 v3.5.1 // indirect + github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc // indirect + github.com/vmihailenco/msgpack/v5 v5.4.1 // indirect + github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect + go.opentelemetry.io/otel v1.37.0 // indirect + go.opentelemetry.io/otel/trace v1.37.0 // indirect golang.org/x/sys v0.37.0 // indirect + mellium.im/sasl v0.3.2 // indirect ) diff --git a/go.sum b/go.sum index 28ab73f..96a9846 100644 --- a/go.sum +++ b/go.sum @@ -1,16 +1,48 @@ github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c= github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= +github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= github.com/oschwald/geoip2-golang v1.13.0 h1:Q44/Ldc703pasJeP5V9+aFSZFmBN7DKHbNsSFzQATJI= github.com/oschwald/geoip2-golang v1.13.0/go.mod h1:P9zG+54KPEFOliZ29i7SeYZ/GM6tfEL+rgSn03hYuUo= github.com/oschwald/maxminddb-golang v1.13.0 h1:R8xBorY71s84yO06NgTmQvqvTvlS/bnYZrrWX1MElnU= github.com/oschwald/maxminddb-golang v1.13.0/go.mod h1:BU0z8BfFVhi1LQaonTwwGQlsHUEu9pWNdMfmq4ztm0o= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/puzpuzpuz/xsync/v3 v3.5.1 h1:GJYJZwO6IdxN/IKbneznS6yPkVC+c3zyY/j19c++5Fg= +github.com/puzpuzpuz/xsync/v3 v3.5.1/go.mod h1:VjzYrABPabuM4KyBh1Ftq6u8nhwY5tBPKP9jpmh0nnA= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc h1:9lRDQMhESg+zvGYmW5DyG0UqvY96Bu5QYsTLvCHdrgo= +github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc/go.mod h1:bciPuU6GHm1iF1pBvUfxfsH0Wmnc2VbpgvbI9ZWuIRs= +github.com/uptrace/bun v1.2.15 h1:Ut68XRBLDgp9qG9QBMa9ELWaZOmzHNdczHQdrOZbEFE= +github.com/uptrace/bun v1.2.15/go.mod h1:Eghz7NonZMiTX/Z6oKYytJ0oaMEJ/eq3kEV4vSqG038= +github.com/uptrace/bun/dialect/pgdialect v1.2.15 h1:er+/3giAIqpfrXJw+KP9B7ujyQIi5XkPnFmgjAVL6bA= +github.com/uptrace/bun/dialect/pgdialect v1.2.15/go.mod h1:QSiz6Qpy9wlGFsfpf7UMSL6mXAL1jDJhFwuOVacCnOQ= +github.com/uptrace/bun/driver/pgdriver v1.2.15 h1:eZZ60ZtUUE6jjv6VAI1pCMaTgtx3sxmChQzwbvchOOo= +github.com/uptrace/bun/driver/pgdriver v1.2.15/go.mod h1:s2zz/BAeScal4KLFDI8PURwATN8s9RDBsElEbnPAjv4= +github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8= +github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok= +github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g= +github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds= +go.opentelemetry.io/otel v1.37.0 h1:9zhNfelUvx0KBfu/gb+ZgeAfAgtWrfHJZcAqFC228wQ= +go.opentelemetry.io/otel v1.37.0/go.mod h1:ehE/umFRLnuLa/vSccNq9oS1ErUlkkK71gMcN34UG8I= +go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mxVK7z4= +go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0= golang.org/x/crypto v0.43.0 h1:dduJYIi3A3KOfdGOHX8AVZ/jGiyPa3IbBozJ5kNuE04= golang.org/x/crypto v0.43.0/go.mod h1:BFbav4mRNlXJL4wNeejLpWxB7wMbc79PdRGhWKncxR0= golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ= golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/term v0.36.0 h1:zMPR+aF8gfksFprF/Nc/rd1wRS1EI6nDBGyWAvDzx2Q= golang.org/x/term v0.36.0/go.mod h1:Qu394IJq6V6dCBRgwqshf3mPF85AqzYEzofzRdZkWss= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +mellium.im/sasl v0.3.2 h1:PT6Xp7ccn9XaXAnJ03FcEjmAn7kK1x7aoXV6F+Vmrl0= +mellium.im/sasl v0.3.2/go.mod h1:NKXDi1zkr+BlMHLQjY3ofYuU4KSPFxknb8mfEu6SveY= diff --git a/inc.env b/inc.env index c776632..4b41664 100644 --- a/inc.env +++ b/inc.env @@ -1,2 +1,3 @@ PORT=2222 -GEOIP_DB=./GeoLite2-City.mmdb \ No newline at end of file +GEOIP_DB=./GeoLite2-City.mmdb +DB_DSN="postgrtesql://postgres:password@localhost/postgres" \ No newline at end of file diff --git a/main.go b/main.go index 2761cbe..931852e 100644 --- a/main.go +++ b/main.go @@ -6,15 +6,17 @@ import ( "slices" "strings" + "sshchat/db" "sshchat/utils" "github.com/gliderlabs/ssh" "github.com/oschwald/geoip2-golang" + "github.com/uptrace/bun" ) var config = utils.GetConfig() -func sessionHandler(s ssh.Session, geoip *geoip2.Reader) { +func sessionHandler(s ssh.Session, geoip *geoip2.Reader, pgDb *bun.DB) { ptyReq, _, isPty := s.Pty() if !isPty { _, _ = fmt.Fprintln(s, "Err: PTY requires. Reconnect with -t option.") @@ -61,11 +63,17 @@ func sessionHandler(s ssh.Session, geoip *geoip2.Reader) { func main() { geoip, err := utils.GetDB(config.Geoip) - port := config.Port if err != nil { log.Fatalf("Geoip db is error: %v", err) } + pgDb, err := db.GetDB(config.PgDsn) + if err != nil { + log.Fatalf("DB Connection error: %v", err) + } + + port := config.Port + keys, err := utils.CheckHostKey() if err != nil { log.Print("Failed to check SSH keys: generate one.\n", err) @@ -83,13 +91,17 @@ func main() { s := &ssh.Server{ Addr: ":" + port, Handler: func(s ssh.Session) { - sessionHandler(s, geoip) + sessionHandler(s, geoip, pgDb) }, } for _, key := range keys { s.AddHostKey(key) } + defer func() { + _ = pgDb.Close() + }() + log.Print("Listening on :" + port) log.Fatal(s.ListenAndServe()) } diff --git a/utils/config.go b/utils/config.go index eee1d7e..47b2e8b 100644 --- a/utils/config.go +++ b/utils/config.go @@ -12,6 +12,7 @@ type Config struct { Port string Geoip string CountryBlacklist []string + PgDsn string } func GetConfig() *Config { @@ -21,12 +22,14 @@ func GetConfig() *Config { } port := os.Getenv("PORT") - geoip_dbfile := os.Getenv("GEOIP_DB") - country_blacklist := os.Getenv("COUNTRY_BLACKLIST") + geoipDbfile := os.Getenv("GEOIP_DB") + countryBlacklist := os.Getenv("COUNTRY_BLACKLIST") + pgDsn := os.Getenv("DB_DSN") return &Config{ Port: port, - Geoip: geoip_dbfile, - CountryBlacklist: strings.Split(country_blacklist, ","), + Geoip: geoipDbfile, + CountryBlacklist: strings.Split(countryBlacklist, ","), + PgDsn: pgDsn, } }