From ad7fc189ccba57256552f2e021d765c57ce010b0 Mon Sep 17 00:00:00 2001
From: Drew DeVault <sir@cmpwn.com>
Date: Mon, 11 Mar 2019 17:09:09 -0400
Subject: [PATCH] Run user tasks in a pty and add ansi rendering

---
 buildsrht/blueprints/jobs.py | 17 ++++++++++++++---
 worker/context.go            |  2 +-
 worker/go.mod                |  1 +
 worker/go.sum                |  2 ++
 worker/tasks.go              | 18 ++++++++++++------
 5 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/buildsrht/blueprints/jobs.py b/buildsrht/blueprints/jobs.py
index 4a98fbb..c143e48 100644
--- a/buildsrht/blueprints/jobs.py
+++ b/buildsrht/blueprints/jobs.py
@@ -1,3 +1,4 @@
+from ansi2html import Ansi2HTMLConverter
 from flask import Blueprint, render_template, request, abort, redirect, session
 from flask import Response
 from flask_login import current_user
@@ -217,7 +218,10 @@ def tag_svg(username, path):
 
 log_max = 131072
 
+ansi = Ansi2HTMLConverter(scheme="mint-terminal")
+
 def logify(text, task, log_url):
+    text = ansi.convert(text, full=False)
     if len(text) >= log_max:
         text = text[-log_max:]
         try:
@@ -231,19 +235,26 @@ def logify(text, task, log_url):
                     f'<a target="_blank" href="{escape(log_url)}">'
                         'Click here to download the full log</a>.'
                     '</span>\n\n')
-                + escape(text)
+                + Markup(text)
                 + Markup('</pre>'))
         linenos = Markup('<pre>\n\n\n')
     else:
         nlines = len(text.splitlines())
-        text = Markup('<pre>') + escape(text) + Markup('</pre>')
+        text = (Markup('<pre>')
+                + Markup(text)
+                + Markup('</pre>'))
         linenos = Markup('<pre>')
     for no in range(1, nlines + 1):
         linenos += Markup(f"<a href='#{escape(task)}-{no-1}'>{no}</a>")
         if no != nlines:
             linenos += Markup("\n")
     linenos += Markup("</pre>")
-    return Markup('<td>') + linenos + Markup('</td><td>') + text + Markup('</td>')
+    return (Markup('<td>')
+            + linenos
+            + Markup('</td><td>')
+            + Markup(ansi.produce_headers())
+            + text
+            + Markup('</td>'))
 
 @jobs.route("/~<username>/job/<int:job_id>")
 def job_by_id(username, job_id):
diff --git a/worker/context.go b/worker/context.go
index bc594ff..191100a 100644
--- a/worker/context.go
+++ b/worker/context.go
@@ -191,7 +191,7 @@ func (ctx *JobContext) Control(
 func (ctx *JobContext) SSH(args ...string) *exec.Cmd {
 	sport := strconv.Itoa(ctx.Port)
 	return exec.CommandContext(ctx.Context, "ssh",
-		append([]string{"-q",
+		append([]string{"-q", "-t",
 			"-p", sport,
 			"-o", "UserKnownHostsFile=/dev/null",
 			"-o", "StrictHostKeyChecking=no",
diff --git a/worker/go.mod b/worker/go.mod
index dc9317a..eb5fbd2 100644
--- a/worker/go.mod
+++ b/worker/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/go-redis/redis v6.14.1+incompatible
 	github.com/golang/protobuf v1.2.0 // indirect
 	github.com/gomodule/redigo v2.0.0+incompatible // indirect
+	github.com/kr/pty v1.1.3
 	github.com/lib/pq v1.0.0
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
 	github.com/mitchellh/mapstructure v1.0.0
diff --git a/worker/go.sum b/worker/go.sum
index 504c34b..c82f4bc 100644
--- a/worker/go.sum
+++ b/worker/go.sum
@@ -8,6 +8,8 @@ github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/gomodule/redigo v2.0.0+incompatible h1:K/R+8tc58AaqLkqG2Ol3Qk+DR/TlNuhuh457pBFPtt0=
 github.com/gomodule/redigo v2.0.0+incompatible/go.mod h1:B4C85qUVwatsJoIUNIfCRsp7qO0iAmpGFZ4EELWSbC4=
+github.com/kr/pty v1.1.3 h1:/Um6a/ZmD5tF7peoOJ5oN5KMQ0DrGVQSXLNwyckutPk=
+github.com/kr/pty v1.1.3/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/lib/pq v1.0.0 h1:X5PMW56eZitiTeO7tKzZxFCSpbFZJtkMMooicw2us9A=
 github.com/lib/pq v1.0.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU=
diff --git a/worker/tasks.go b/worker/tasks.go
index 7791548..353a962 100644
--- a/worker/tasks.go
+++ b/worker/tasks.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/url"
 	"os"
@@ -14,6 +15,7 @@ import (
 	"time"
 
 	"github.com/go-redis/redis"
+	"github.com/kr/pty"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -98,7 +100,7 @@ func (ctx *JobContext) Settle() error {
 	go func() {
 		for {
 			attempt++
-			check := ctx.SSH("echo", "hello world")
+			check := ctx.SSH("printf", "'hello world'")
 			pipe, _ := check.StdoutPipe()
 			if err := check.Start(); err != nil {
 				done <- err
@@ -106,7 +108,7 @@ func (ctx *JobContext) Settle() error {
 			}
 			stdout, _ := ioutil.ReadAll(pipe)
 			if err := check.Wait(); err == nil {
-				if string(stdout) == "hello world\n" {
+				if string(stdout) == "hello world" {
 					done <- nil
 					return
 				} else {
@@ -379,6 +381,7 @@ func (ctx *JobContext) RunTasks() error {
 			logfd *os.File
 			name  string
 			ssh   *exec.Cmd
+			tty   *os.File
 		)
 		for name, _ = range task {
 			break
@@ -393,14 +396,17 @@ func (ctx *JobContext) RunTasks() error {
 
 		ssh = ctx.SSH(path.Join(".", ".tasks", name))
 		if logfd, err = os.Create(path.Join(ctx.LogDir, name, "log")); err != nil {
-
 			err = errors.Wrap(err, "Creating log file")
 			goto fail
 		}
-		ssh.Stdout = logfd
-		ssh.Stderr = logfd
+		tty, err = pty.Start(ssh)
+		if err != nil {
+			err = errors.Wrap(err, "Allocating pty")
+			goto fail
+		}
+		go io.Copy(logfd, tty)
 
-		if err = ssh.Run(); err != nil {
+		if err = ssh.Wait(); err != nil {
 			exiterr, ok := err.(*exec.ExitError)
 			if !ok {
 				goto fail
-- 
2.38.5