From b14477bfeb81c422eae5d0c542fd579805cdf5aa Mon Sep 17 00:00:00 2001
From: Lucas TESSON <lucastesson@protonmail.com>
Date: Mon, 15 Apr 2024 23:40:33 +0200
Subject: [PATCH] fix: login/logout session cookie handling

---
 api/login.go      | 10 ++++++++++
 api/logout.go     | 11 +++++++++++
 api/setup_test.go |  3 ++-
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/api/login.go b/api/login.go
index b5f0f3b..20c6005 100644
--- a/api/login.go
+++ b/api/login.go
@@ -43,5 +43,15 @@ func (client *Client) Login(params *LoginParams, opts ...Option) error {
 		return err
 	}
 	client.nonce = nonce
+
+	u, _ := url.Parse(client.url)
+	hds := client.sub.Jar.Cookies(u)
+	for _, hd := range hds {
+		if hd.Name == "session" {
+			client.session = hd.Value
+			break
+		}
+	}
+
 	return nil
 }
diff --git a/api/logout.go b/api/logout.go
index caa64f7..62c1ed3 100644
--- a/api/logout.go
+++ b/api/logout.go
@@ -3,6 +3,7 @@ package api
 import (
 	"fmt"
 	"net/http"
+	"net/url"
 )
 
 func (client *Client) Logout(opts ...Option) error {
@@ -23,5 +24,15 @@ func (client *Client) Logout(opts ...Option) error {
 		return err
 	}
 	client.nonce = nonce
+
+	u, _ := url.Parse(client.url)
+	hds := client.sub.Jar.Cookies(u)
+	for _, hd := range hds {
+		if hd.Name == "session" {
+			client.session = hd.Value
+			break
+		}
+	}
+
 	return nil
 }
diff --git a/api/setup_test.go b/api/setup_test.go
index 36fbe15..0485aaf 100644
--- a/api/setup_test.go
+++ b/api/setup_test.go
@@ -81,7 +81,7 @@ func Test_F_Setup(t *testing.T) {
 		return
 	}
 
-	// 1e. Relog :)
+	// 1e. Relog, and drop the use of the API token because... why not ¯\_(ツ)_/¯
 	err = client.Login(&api.LoginParams{
 		Name:     "ctfer",
 		Password: "password",
@@ -89,6 +89,7 @@ func Test_F_Setup(t *testing.T) {
 	if !assert.Nil(err, "got error: %s", err) {
 		return
 	}
+	client.SetAPIKey("")
 
 	// 2. Create a challenge
 	chall, err := client.PostChallenges(&api.PostChallengesParams{