<commit_before>//Copyright 2014 Rana Ian. All rights reserved.
//Use of this source code is governed by The MIT License
//found in the accompanying LICENSE file.

package ora_test

import (
	"database/sql"
	"fmt"
	"os"
	"testing"

	"gopkg.in/rana/ora.v3"
)

func Test_open_cursors_db(t *testing.T) {
	//enableLogging(t)
	// This needs "GRANT SELECT ANY DICTIONARY TO test"
	// or at least "GRANT SELECT ON v_$mystat TO test".
	// use 'opened cursors current' statistic#=5 to determine opened cursors on oracle server
	// SELECT A.STATISTIC#, A.NAME, B.VALUE
	// FROM V$STATNAME A, V$MYSTAT B
	// WHERE A.STATISTIC# = B.STATISTIC#
	stmt, err := testDb.Prepare("SELECT VALUE FROM V$MYSTAT WHERE STATISTIC#=5")
	if err != nil {
		t.Fatal(err)
	}
	var before, after int
	if err = stmt.QueryRow().Scan(&before); err != nil {
		t.Skip(err)
	}
	rounds := 100
	for i := 0; i < rounds; i++ {
		func() {
			stmt, err := testDb.Prepare("SELECT 1 FROM user_objects WHERE ROWNUM < 100")
			if err != nil {
				t.Fatal(err)
			}
			defer stmt.Close()
			rows, err := stmt.Query()
			if err != nil {
				t.Errorf("SELECT: %v", err)
				return
			}
			defer rows.Close()
			j := 0
			for rows.Next() {
				j++
			}
			//t.Logf("%d objects, error=%v", j, rows.Err())
		}()
	}
	if err = stmt.QueryRow().Scan(&after); err != nil {
		t.Fatal(err)
	}
	if after-before >= rounds {
		t.Errorf("before=%d after=%d, awaited less than %d increment!", before, after, rounds)
		return
	}
	t.Logf("before=%d after=%d", before, after)
}

func TestSelectNull_db(t *testing.T) {
	ora.Cfg().Log.Rset.BeginRow = true
	//enableLogging(t)
	var (
		s   string
		oS  ora.String
		i   int64
		oI  ora.Int64
		tim ora.Time
	)
	for tN, tC := range []struct {
		Field string
		Dest  interface{}
	}{
		{"''", &s},
		{"''", &oS},
		{"NULL + 0", &i},
		{"NULL + 0", &oI},
		{"SYSDATE + NULL", &tim},
	} {
		qry := "SELECT " + tC.Field + " x FROM DUAL"
		rows, err := testDb.Query(qry)
		if err != nil {
			t.Errorf("%d. %s: %v", tN, qry, err)
			return
		}
		for rows.Next() {
			if err = rows.Scan(&tC.Dest); err != nil {
				t.Errorf("%d. Scan: %v", tN, err)
				break
			}
		}
		if rows.Err() != nil {
			t.Errorf("%d. rows: %v", tN, rows.Err())
		}
		rows.Close()
	}
}

func Test_numberP38S0Identity_db(t *testing.T) {
	tableName := tableName()
	stmt, err := testDb.Prepare(createTableSql(tableName, 1, numberP38S0Identity, varchar2C48))
	if err == nil {
		defer stmt.Close()
		_, err = stmt.Exec()
	}
	if err != nil {
		t.Skipf("SKIP create table with identity: %v", err)
		return
	}
	defer dropTableDB(testDb, t, tableName)

	stmt, err = testDb.Prepare(fmt.Sprintf("insert into %v (c2) values ('go') returning c1 /*lastInsertId*/ into :c1", tableName))
	defer stmt.Close()

	// pass nil to Exec when using 'returning into' clause with sql.DB
	result, err := stmt.Exec(nil)
	testErr(err, t)
	actual, err := result.LastInsertId()
	testErr(err, t)
	if 1 != actual {
		t.Fatalf("LastInsertId: expected(%v), actual(%v)", 1, actual)
	}
}

func Test_numberP38S0_int64_db(t *testing.T) {
	testBindDefineDB(gen_int64(), t, numberP38S0)
}

func Test_numberP38S0Null_int64_db(t *testing.T) {
	testBindDefineDB(gen_int64(), t, numberP38S0Null)
}

func Test_numberP16S15_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, numberP16S15)
}

func Test_numberP16S15Null_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, numberP16S15Null)
}

func Test_binaryDouble_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryDouble)
}

func Test_binaryDoubleNull_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryDoubleNull)
}

func Test_binaryFloat_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryFloat)
}

func Test_binaryFloatNull_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryFloatNull)
}

func Test_floatP126_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, floatP126)
}

func Test_floatP126Null_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, floatP126Null)
}

func Test_date_time_db(t *testing.T) {
	testBindDefineDB(gen_date(), t, dateNotNull)
}

func Test_dateNull_time_db(t *testing.T) {
	testBindDefineDB(gen_date(), t, dateNull)
}

func Test_timestampP9_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampP9)
}

func Test_timestampP9Null_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampP9Null)
}

func Test_timestampTzP9_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampTzP9)
}

func Test_timestampTzP9Null_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampTzP9Null)
}

func Test_timestampLtzP9_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampLtzP9)
}

func Test_timestampLtzP9Null_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampLtzP9Null)
}

func Test_charB48_string_db(t *testing.T) {
	enableLogging(t)
	testBindDefineDB(gen_string48(), t, charB48)
}

func Test_charB48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, charB48Null)
}

func Test_charC48_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, charC48)
}

func Test_charC48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, charC48Null)
}

func Test_nchar48_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, nchar48)
}

func Test_nchar48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, nchar48Null)
}

func Test_varcharB48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharB48)
}

func Test_varcharB48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharB48Null)
}

func Test_varcharC48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharC48)
}

func Test_varcharC48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharC48Null)
}

func Test_varchar2B48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2B48)
}

func Test_varchar2B48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2B48Null)
}

func Test_varchar2C48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2C48)
}

func Test_varchar2C48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2C48Null)
}

func Test_nvarchar248_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nvarchar248)
}

func Test_nvarchar248Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nvarchar248Null)
}

func Test_long_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, long)
}

func Test_longNull_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, longNull)
}

func Test_clob_string_db(t *testing.T) {
	//enableLogging(t)
	testBindDefineDB(gen_string(), t, clob)
}

func Test_clobNull_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, clobNull)
}

func Test_nclob_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nclob)
}

func Test_nclobNull_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nclobNull)
}

func Test_charB1_bool_true_db(t *testing.T) {
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charB1)
}

func Test_charB1Null_bool_true_db(t *testing.T) {
	//enableLogging(t)
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charB1Null)
}

func Test_charC1_bool_true_db(t *testing.T) {
	//enableLogging(t)
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charC1)
}

func Test_charC1Null_bool_true_db(t *testing.T) {
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charC1Null)
}

func Test_longRaw_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, longRaw)
}

func Test_longRawNull_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, longRawNull)
}

func Test_raw2000_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(2000), t, raw2000)
}

func Test_raw2000Null_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(2000), t, raw2000Null)
}

func Test_blob_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, blob)
}

func Test_blobNull_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, blobNull)
}

func TestSysdba(t *testing.T) {
	u := os.Getenv("GO_ORA_DRV_TEST_SYSDBA_USERNAME")
	p := os.Getenv("GO_ORA_DRV_TEST_SYSDBA_PASSWORD")
	if u == "" {
		u = testSesCfg.Username
		p = testSesCfg.Password
	}
	dsn := fmt.Sprintf("%s/%s@%s AS SYSDBA", u, p, testSrvCfg.Dblink)
	db, err := sql.Open("ora", dsn)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if err := db.Ping(); err != nil {
		t.Skipf("%q: %v", dsn, err)
	}
}

func TestZeroRowsAffected(t *testing.T) {
	tableName := tableName()
	if _, err := testDb.Exec("CREATE TABLE " + tableName + " (id NUMBER(3))"); err != nil {
		t.Fatal(err)
	}
	defer testDb.Exec("DROP TABLE " + tableName)
	res, err := testDb.Exec("UPDATE " + tableName + " SET id=1 WHERE 1=0")
	if err != nil {
		t.Fatal(err)
	}
	if ra, err := res.RowsAffected(); err != nil {
		t.Error(err)
	} else if ra != 0 {
		t.Errorf("got %d, wanted 0 rows affected!")
	}
	if _, err := res.LastInsertId(); err == nil {
		t.Error("wanted error for LastInsertId, got nil")
	}
}
<commit_msg>add TestSetMaxConnLifetime<commit_after>//Copyright 2014 Rana Ian. All rights reserved.
//Use of this source code is governed by The MIT License
//found in the accompanying LICENSE file.

package ora_test

import (
	"database/sql"
	"fmt"
	"math/rand"
	"os"
	"sync"
	"testing"
	"time"

	"gopkg.in/rana/ora.v3"
)

func Test_open_cursors_db(t *testing.T) {
	//enableLogging(t)
	// This needs "GRANT SELECT ANY DICTIONARY TO test"
	// or at least "GRANT SELECT ON v_$mystat TO test".
	// use 'opened cursors current' statistic#=5 to determine opened cursors on oracle server
	// SELECT A.STATISTIC#, A.NAME, B.VALUE
	// FROM V$STATNAME A, V$MYSTAT B
	// WHERE A.STATISTIC# = B.STATISTIC#
	stmt, err := testDb.Prepare("SELECT VALUE FROM V$MYSTAT WHERE STATISTIC#=5")
	if err != nil {
		t.Fatal(err)
	}
	var before, after int
	if err = stmt.QueryRow().Scan(&before); err != nil {
		t.Skip(err)
	}
	rounds := 100
	for i := 0; i < rounds; i++ {
		func() {
			stmt, err := testDb.Prepare("SELECT 1 FROM user_objects WHERE ROWNUM < 100")
			if err != nil {
				t.Fatal(err)
			}
			defer stmt.Close()
			rows, err := stmt.Query()
			if err != nil {
				t.Errorf("SELECT: %v", err)
				return
			}
			defer rows.Close()
			j := 0
			for rows.Next() {
				j++
			}
			//t.Logf("%d objects, error=%v", j, rows.Err())
		}()
	}
	if err = stmt.QueryRow().Scan(&after); err != nil {
		t.Fatal(err)
	}
	if after-before >= rounds {
		t.Errorf("before=%d after=%d, awaited less than %d increment!", before, after, rounds)
		return
	}
	t.Logf("before=%d after=%d", before, after)
}

func TestSelectNull_db(t *testing.T) {
	ora.Cfg().Log.Rset.BeginRow = true
	//enableLogging(t)
	var (
		s   string
		oS  ora.String
		i   int64
		oI  ora.Int64
		tim ora.Time
	)
	for tN, tC := range []struct {
		Field string
		Dest  interface{}
	}{
		{"''", &s},
		{"''", &oS},
		{"NULL + 0", &i},
		{"NULL + 0", &oI},
		{"SYSDATE + NULL", &tim},
	} {
		qry := "SELECT " + tC.Field + " x FROM DUAL"
		rows, err := testDb.Query(qry)
		if err != nil {
			t.Errorf("%d. %s: %v", tN, qry, err)
			return
		}
		for rows.Next() {
			if err = rows.Scan(&tC.Dest); err != nil {
				t.Errorf("%d. Scan: %v", tN, err)
				break
			}
		}
		if rows.Err() != nil {
			t.Errorf("%d. rows: %v", tN, rows.Err())
		}
		rows.Close()
	}
}

func TestSetConnMaxLifetime(t *testing.T) {
	var db *sql.DB
	var err error
	db, err = sql.Open("ora", testConStr)
	if err != nil {
		t.Fatal(err)
	}
	db.SetMaxIdleConns(2)
	db.SetConnMaxLifetime(1 * time.Second)
	defer db.Close()

	done := make(chan struct{})
	var wg sync.WaitGroup
	dbRoutine := func() {
		defer wg.Done()
		for {
			select {
			case <-done:
				return
			default:
			}
			var temp int
			db.QueryRow("SELECT 1 FROM DUAL").Scan(&temp)
			if rand.Int()%10 == 0 {
				time.Sleep(50 * time.Millisecond)
			}
		}
	}

	for i := 0; i < 8; i++ {
		wg.Add(1)
		go dbRoutine()
	}
	time.Sleep(8 * time.Second)
	close(done)
	wg.Wait()
}

func Test_numberP38S0Identity_db(t *testing.T) {
	tableName := tableName()
	stmt, err := testDb.Prepare(createTableSql(tableName, 1, numberP38S0Identity, varchar2C48))
	if err == nil {
		defer stmt.Close()
		_, err = stmt.Exec()
	}
	if err != nil {
		t.Skipf("SKIP create table with identity: %v", err)
		return
	}
	defer dropTableDB(testDb, t, tableName)

	stmt, err = testDb.Prepare(fmt.Sprintf("insert into %v (c2) values ('go') returning c1 /*lastInsertId*/ into :c1", tableName))
	defer stmt.Close()

	// pass nil to Exec when using 'returning into' clause with sql.DB
	result, err := stmt.Exec(nil)
	testErr(err, t)
	actual, err := result.LastInsertId()
	testErr(err, t)
	if 1 != actual {
		t.Fatalf("LastInsertId: expected(%v), actual(%v)", 1, actual)
	}
}

func Test_numberP38S0_int64_db(t *testing.T) {
	testBindDefineDB(gen_int64(), t, numberP38S0)
}

func Test_numberP38S0Null_int64_db(t *testing.T) {
	testBindDefineDB(gen_int64(), t, numberP38S0Null)
}

func Test_numberP16S15_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, numberP16S15)
}

func Test_numberP16S15Null_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, numberP16S15Null)
}

func Test_binaryDouble_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryDouble)
}

func Test_binaryDoubleNull_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryDoubleNull)
}

func Test_binaryFloat_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryFloat)
}

func Test_binaryFloatNull_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, binaryFloatNull)
}

func Test_floatP126_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, floatP126)
}

func Test_floatP126Null_float64_db(t *testing.T) {
	testBindDefineDB(gen_float64(), t, floatP126Null)
}

func Test_date_time_db(t *testing.T) {
	testBindDefineDB(gen_date(), t, dateNotNull)
}

func Test_dateNull_time_db(t *testing.T) {
	testBindDefineDB(gen_date(), t, dateNull)
}

func Test_timestampP9_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampP9)
}

func Test_timestampP9Null_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampP9Null)
}

func Test_timestampTzP9_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampTzP9)
}

func Test_timestampTzP9Null_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampTzP9Null)
}

func Test_timestampLtzP9_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampLtzP9)
}

func Test_timestampLtzP9Null_time_db(t *testing.T) {
	testBindDefineDB(gen_time(), t, timestampLtzP9Null)
}

func Test_charB48_string_db(t *testing.T) {
	enableLogging(t)
	testBindDefineDB(gen_string48(), t, charB48)
}

func Test_charB48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, charB48Null)
}

func Test_charC48_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, charC48)
}

func Test_charC48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, charC48Null)
}

func Test_nchar48_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, nchar48)
}

func Test_nchar48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string48(), t, nchar48Null)
}

func Test_varcharB48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharB48)
}

func Test_varcharB48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharB48Null)
}

func Test_varcharC48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharC48)
}

func Test_varcharC48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varcharC48Null)
}

func Test_varchar2B48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2B48)
}

func Test_varchar2B48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2B48Null)
}

func Test_varchar2C48_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2C48)
}

func Test_varchar2C48Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, varchar2C48Null)
}

func Test_nvarchar248_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nvarchar248)
}

func Test_nvarchar248Null_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nvarchar248Null)
}

func Test_long_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, long)
}

func Test_longNull_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, longNull)
}

func Test_clob_string_db(t *testing.T) {
	//enableLogging(t)
	testBindDefineDB(gen_string(), t, clob)
}

func Test_clobNull_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, clobNull)
}

func Test_nclob_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nclob)
}

func Test_nclobNull_string_db(t *testing.T) {
	testBindDefineDB(gen_string(), t, nclobNull)
}

func Test_charB1_bool_true_db(t *testing.T) {
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charB1)
}

func Test_charB1Null_bool_true_db(t *testing.T) {
	//enableLogging(t)
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charB1Null)
}

func Test_charC1_bool_true_db(t *testing.T) {
	//enableLogging(t)
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charC1)
}

func Test_charC1Null_bool_true_db(t *testing.T) {
	defer setC1Bool()()
	testBindDefineDB(gen_boolTrue(), t, charC1Null)
}

func Test_longRaw_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, longRaw)
}

func Test_longRawNull_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, longRawNull)
}

func Test_raw2000_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(2000), t, raw2000)
}

func Test_raw2000Null_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(2000), t, raw2000Null)
}

func Test_blob_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, blob)
}

func Test_blobNull_bytes_db(t *testing.T) {
	testBindDefineDB(gen_bytes(9), t, blobNull)
}

func TestSysdba(t *testing.T) {
	u := os.Getenv("GO_ORA_DRV_TEST_SYSDBA_USERNAME")
	p := os.Getenv("GO_ORA_DRV_TEST_SYSDBA_PASSWORD")
	if u == "" {
		u = testSesCfg.Username
		p = testSesCfg.Password
	}
	dsn := fmt.Sprintf("%s/%s@%s AS SYSDBA", u, p, testSrvCfg.Dblink)
	db, err := sql.Open("ora", dsn)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if err := db.Ping(); err != nil {
		t.Skipf("%q: %v", dsn, err)
	}
}

func TestZeroRowsAffected(t *testing.T) {
	tableName := tableName()
	if _, err := testDb.Exec("CREATE TABLE " + tableName + " (id NUMBER(3))"); err != nil {
		t.Fatal(err)
	}
	defer testDb.Exec("DROP TABLE " + tableName)
	res, err := testDb.Exec("UPDATE " + tableName + " SET id=1 WHERE 1=0")
	if err != nil {
		t.Fatal(err)
	}
	if ra, err := res.RowsAffected(); err != nil {
		t.Error(err)
	} else if ra != 0 {
		t.Errorf("got %d, wanted 0 rows affected!")
	}
	if _, err := res.LastInsertId(); err == nil {
		t.Error("wanted error for LastInsertId, got nil")
	}
}
<|endoftext|>


<commit_before>/*
 * ZDNS Copyright 2016 Regents of the University of Michigan
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package main

import (
	"flag"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
	"github.com/zmap/zdns"
	_ "github.com/zmap/zdns/modules/alookup"
	_ "github.com/zmap/zdns/modules/axfr"
	_ "github.com/zmap/zdns/modules/dmarc"
	_ "github.com/zmap/zdns/modules/miekg"
	_ "github.com/zmap/zdns/modules/mxlookup"
	_ "github.com/zmap/zdns/modules/nslookup"
	_ "github.com/zmap/zdns/modules/spf"

	_ "github.com/zmap/zdns/iohandlers/file"
)

func main() {

	var gc zdns.GlobalConf
	// global flags relevant to every lookup module
	flags := flag.NewFlagSet("flags", flag.ExitOnError)
	flags.IntVar(&gc.Threads, "threads", 1000, "number of lightweight go threads")
	flags.IntVar(&gc.GoMaxProcs, "go-processes", 0, "number of OS processes (GOMAXPROCS)")
	flags.StringVar(&gc.NamePrefix, "prefix", "", "name to be prepended to what's passed in (e.g., www.)")
	flags.BoolVar(&gc.AlexaFormat, "alexa", false, "is input file from Alexa Top Million download")
	flags.BoolVar(&gc.IterativeResolution, "iterative", false, "Perform own iteration instead of relying on recursive resolver")
	flags.StringVar(&gc.InputFilePath, "input-file", "-", "names to read")
	flags.StringVar(&gc.OutputFilePath, "output-file", "-", "where should JSON output be saved")
	flags.StringVar(&gc.MetadataFilePath, "metadata-file", "", "where should JSON metadata be saved")
	flags.StringVar(&gc.LogFilePath, "log-file", "", "where should JSON logs be saved")

	flags.StringVar(&gc.ResultVerbosity, "result-verbosity", "normal", "Sets verbosity of each output record. Options: short, normal, long, trace")
	flags.StringVar(&gc.IncludeInOutput, "include-fields", "", "Comma separated list of fields to additionally output beyond result verbosity. Options: class, protocol, ttl, resolver, flags")

	flags.IntVar(&gc.Verbosity, "verbosity", 3, "log verbosity: 1 (lowest)--5 (highest)")
	flags.IntVar(&gc.Retries, "retries", 1, "how many times should zdns retry query if timeout or temporary failure")
	flags.IntVar(&gc.MaxDepth, "max-depth", 10, "how deep should we recurse when performing iterative lookups")
	flags.IntVar(&gc.CacheSize, "cache-size", 10000, "how many items can be stored in internal recursive cache")
	flags.StringVar(&gc.InputHandler, "input-handler", "file", "handler to input names")
	flags.StringVar(&gc.OutputHandler, "output-handler", "file", "handler to output names")
	flags.BoolVar(&gc.TCPOnly, "tcp-only", false, "Only perform lookups over TCP")
	flags.BoolVar(&gc.UDPOnly, "udp-only", false, "Only perform lookups over UDP")
	servers_string := flags.String("name-servers", "", "comma-delimited list of DNS servers to use. If no port is specified, defaults to 53.")
	config_file := flags.String("conf-file", "/etc/resolv.conf", "config file for DNS servers")
	timeout := flags.Int("timeout", 15, "timeout for resolving an individual name")
	iterationTimeout := flags.Int("iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
	class_string := flags.String("class", "INET", "DNS class to query. Options: INET, CSNET, CHAOS, HESIOD, NONE, ANY. Default: INET.")
	nanoSeconds := flags.Bool("nanoseconds", false, "Use nanosecond resolution timestamps")
	// allow module to initialize and add its own flags before we parse
	if len(os.Args) < 2 {
		log.Fatal("No lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	gc.Module = strings.ToUpper(os.Args[1])
	factory := zdns.GetLookup(gc.Module)
	if factory == nil {
		flags.Parse(os.Args[1:])
		log.Fatal("Invalid lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	factory.AddFlags(flags)
	flags.Parse(os.Args[2:])
	// Do some basic sanity checking
	// setup global logging
	if gc.LogFilePath != "" {
		f, err := os.OpenFile(gc.LogFilePath, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			log.Fatalf("Unable to open log file (%s): %s", gc.LogFilePath, err.Error())
		}
		log.SetOutput(f)
	}
	// Translate the assigned verbosity level to a logrus log level.
	switch gc.Verbosity {
	case 1: // Fatal
		log.SetLevel(log.FatalLevel)
	case 2: // Error
		log.SetLevel(log.ErrorLevel)
	case 3: // Warnings  (default)
		log.SetLevel(log.WarnLevel)
	case 4: // Information
		log.SetLevel(log.InfoLevel)
	case 5: // Debugging
		log.SetLevel(log.DebugLevel)
	default:
		log.Fatal("Unknown verbosity level specified. Must be between 1 (lowest)--5 (highest)")
	}
	// complete post facto global initialization based on command line arguments
	gc.Timeout = time.Duration(time.Second * time.Duration(*timeout))
	gc.IterationTimeout = time.Duration(time.Second * time.Duration(*iterationTimeout))
	// class initialization
	switch strings.ToUpper(*class_string) {
	case "INET", "IN":
		gc.Class = dns.ClassINET
	case "CSNET", "CS":
		gc.Class = dns.ClassCSNET
	case "CHAOS", "CH":
		gc.Class = dns.ClassCHAOS
	case "HESIOD", "HS":
		gc.Class = dns.ClassHESIOD
	case "NONE":
		gc.Class = dns.ClassNONE
	case "ANY":
		gc.Class = dns.ClassANY
	default:
		log.Fatal("Unknown record class specified. Valid valued are INET (default), CSNET, CHAOS, HESIOD, NONE, ANY")

	}
	if *servers_string == "" {
		// if we're doing recursive resolution, figure out default OS name servers
		// otherwise, use the set of 13 root name servers
		if gc.IterativeResolution {
			gc.NameServers = zdns.RootServers[:]
		} else {
			ns, err := zdns.GetDNSServers(*config_file)
			if err != nil {
				log.Fatal("Unable to fetch correct name servers:", err.Error())
			}
			gc.NameServers = ns
		}
		gc.NameServersSpecified = false
		log.Info("no name servers specified. will use: ", strings.Join(gc.NameServers, ", "))
	} else {
		ns := strings.Split(*servers_string, ",")
		for i, s := range ns {
			if !strings.Contains(s, ":") {
				ns[i] = s + ":53"
			}
		}
		gc.NameServers = ns
		gc.NameServersSpecified = true
	}
	if *nanoSeconds {
		gc.TimeFormat = time.RFC3339Nano
	} else {
		gc.TimeFormat = time.RFC3339
	}
	if gc.GoMaxProcs < 0 {
		log.Fatal("Invalid argument for --go-processes. Must be >1.")
	}
	if gc.GoMaxProcs != 0 {
		runtime.GOMAXPROCS(gc.GoMaxProcs)
	}
	if gc.UDPOnly && gc.TCPOnly {
		log.Fatal("TCP Only and UDP Only are conflicting")
	}
	// Output Groups are defined by a base + any additional fields that the user wants
	groups := strings.Split(gc.IncludeInOutput, ",")
	if gc.ResultVerbosity != "short" && gc.ResultVerbosity != "normal" && gc.ResultVerbosity != "long" && gc.ResultVerbosity != "trace" {
		log.Fatal("Invalid result verbosity. Options: short, normal, long, trace")
	}

	gc.OutputGroups = append(gc.OutputGroups, gc.ResultVerbosity)
	gc.OutputGroups = append(gc.OutputGroups, groups...)

	if len(flags.Args()) > 0 {
		stat, _ := os.Stdin.Stat()
		// If stdin is piped from the terminal, and we havent specified a file, and if we have unparsed args
		// use them for a dig like reslution
		if (stat.Mode()&os.ModeCharDevice) != 0 && gc.InputFilePath == "-" && len(flags.Args()) == 1 {
			gc.PassedName = flags.Args()[0]
		} else {
			log.Fatal("Unused command line flags: ", flags.Args())
		}
	}

	// some modules require multiple passes over a file (this is really just the case for zone files)
	if !factory.AllowStdIn() && gc.InputFilePath == "-" {
		log.Fatal("Specified module does not allow reading from stdin")
	}

	// allow the factory to initialize itself
	if err := factory.Initialize(&gc); err != nil {
		log.Fatal("Factory was unable to initialize:", err.Error())
	}
	// run it.
	if err := zdns.DoLookups(&factory, &gc); err != nil {
		log.Fatal("Unable to run lookups:", err.Error())
	}
	// allow the factory to initialize itself
	if err := factory.Finalize(); err != nil {
		log.Fatal("Factory was unable to finalize:", err.Error())
	}
}
<commit_msg>Add option which passes resolvers using file #185 (#187)<commit_after>/*
 * ZDNS Copyright 2016 Regents of the University of Michigan
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package main

import (
	"flag"
	"os"
	"runtime"
	"strings"
	"time"
	"io/ioutil"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
	"github.com/zmap/zdns"
	_ "github.com/zmap/zdns/modules/alookup"
	_ "github.com/zmap/zdns/modules/axfr"
	_ "github.com/zmap/zdns/modules/dmarc"
	_ "github.com/zmap/zdns/modules/miekg"
	_ "github.com/zmap/zdns/modules/mxlookup"
	_ "github.com/zmap/zdns/modules/nslookup"
	_ "github.com/zmap/zdns/modules/spf"

	_ "github.com/zmap/zdns/iohandlers/file"
)

func main() {

	var gc zdns.GlobalConf
	// global flags relevant to every lookup module
	flags := flag.NewFlagSet("flags", flag.ExitOnError)
	flags.IntVar(&gc.Threads, "threads", 1000, "number of lightweight go threads")
	flags.IntVar(&gc.GoMaxProcs, "go-processes", 0, "number of OS processes (GOMAXPROCS)")
	flags.StringVar(&gc.NamePrefix, "prefix", "", "name to be prepended to what's passed in (e.g., www.)")
	flags.BoolVar(&gc.AlexaFormat, "alexa", false, "is input file from Alexa Top Million download")
	flags.BoolVar(&gc.IterativeResolution, "iterative", false, "Perform own iteration instead of relying on recursive resolver")
	flags.StringVar(&gc.InputFilePath, "input-file", "-", "names to read")
	flags.StringVar(&gc.OutputFilePath, "output-file", "-", "where should JSON output be saved")
	flags.StringVar(&gc.MetadataFilePath, "metadata-file", "", "where should JSON metadata be saved")
	flags.StringVar(&gc.LogFilePath, "log-file", "", "where should JSON logs be saved")

	flags.StringVar(&gc.ResultVerbosity, "result-verbosity", "normal", "Sets verbosity of each output record. Options: short, normal, long, trace")
	flags.StringVar(&gc.IncludeInOutput, "include-fields", "", "Comma separated list of fields to additionally output beyond result verbosity. Options: class, protocol, ttl, resolver, flags")

	flags.IntVar(&gc.Verbosity, "verbosity", 3, "log verbosity: 1 (lowest)--5 (highest)")
	flags.IntVar(&gc.Retries, "retries", 1, "how many times should zdns retry query if timeout or temporary failure")
	flags.IntVar(&gc.MaxDepth, "max-depth", 10, "how deep should we recurse when performing iterative lookups")
	flags.IntVar(&gc.CacheSize, "cache-size", 10000, "how many items can be stored in internal recursive cache")
	flags.StringVar(&gc.InputHandler, "input-handler", "file", "handler to input names")
	flags.StringVar(&gc.OutputHandler, "output-handler", "file", "handler to output names")
	flags.BoolVar(&gc.TCPOnly, "tcp-only", false, "Only perform lookups over TCP")
	flags.BoolVar(&gc.UDPOnly, "udp-only", false, "Only perform lookups over UDP")
	servers_string := flags.String("name-servers", "", "List of DNS servers to use. Can be passed as comma-delimited string or via @/path/to/file. If no port is specified, defaults to 53.")
	config_file := flags.String("conf-file", "/etc/resolv.conf", "config file for DNS servers")
	timeout := flags.Int("timeout", 15, "timeout for resolving an individual name")
	iterationTimeout := flags.Int("iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
	class_string := flags.String("class", "INET", "DNS class to query. Options: INET, CSNET, CHAOS, HESIOD, NONE, ANY. Default: INET.")
	nanoSeconds := flags.Bool("nanoseconds", false, "Use nanosecond resolution timestamps")
	// allow module to initialize and add its own flags before we parse
	if len(os.Args) < 2 {
		log.Fatal("No lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	gc.Module = strings.ToUpper(os.Args[1])
	factory := zdns.GetLookup(gc.Module)
	if factory == nil {
		flags.Parse(os.Args[1:])
		log.Fatal("Invalid lookup module specified. Valid modules: ", zdns.ValidlookupsString())
	}
	factory.AddFlags(flags)
	flags.Parse(os.Args[2:])
	// Do some basic sanity checking
	// setup global logging
	if gc.LogFilePath != "" {
		f, err := os.OpenFile(gc.LogFilePath, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			log.Fatalf("Unable to open log file (%s): %s", gc.LogFilePath, err.Error())
		}
		log.SetOutput(f)
	}
	// Translate the assigned verbosity level to a logrus log level.
	switch gc.Verbosity {
	case 1: // Fatal
		log.SetLevel(log.FatalLevel)
	case 2: // Error
		log.SetLevel(log.ErrorLevel)
	case 3: // Warnings  (default)
		log.SetLevel(log.WarnLevel)
	case 4: // Information
		log.SetLevel(log.InfoLevel)
	case 5: // Debugging
		log.SetLevel(log.DebugLevel)
	default:
		log.Fatal("Unknown verbosity level specified. Must be between 1 (lowest)--5 (highest)")
	}
	// complete post facto global initialization based on command line arguments
	gc.Timeout = time.Duration(time.Second * time.Duration(*timeout))
	gc.IterationTimeout = time.Duration(time.Second * time.Duration(*iterationTimeout))
	// class initialization
	switch strings.ToUpper(*class_string) {
	case "INET", "IN":
		gc.Class = dns.ClassINET
	case "CSNET", "CS":
		gc.Class = dns.ClassCSNET
	case "CHAOS", "CH":
		gc.Class = dns.ClassCHAOS
	case "HESIOD", "HS":
		gc.Class = dns.ClassHESIOD
	case "NONE":
		gc.Class = dns.ClassNONE
	case "ANY":
		gc.Class = dns.ClassANY
	default:
		log.Fatal("Unknown record class specified. Valid valued are INET (default), CSNET, CHAOS, HESIOD, NONE, ANY")

	}
	if *servers_string == "" {
		// if we're doing recursive resolution, figure out default OS name servers
		// otherwise, use the set of 13 root name servers
		if gc.IterativeResolution {
			gc.NameServers = zdns.RootServers[:]
		} else {
			ns, err := zdns.GetDNSServers(*config_file)
			if err != nil {
				log.Fatal("Unable to fetch correct name servers:", err.Error())
			}
			gc.NameServers = ns
		}
		gc.NameServersSpecified = false
		log.Info("no name servers specified. will use: ", strings.Join(gc.NameServers, ", "))
	} else {
		var ns []string
		if (*servers_string)[0] == '@' {
			filepath := (*servers_string)[1:]
			f, err := ioutil.ReadFile(filepath)
			if err != nil {
				log.Fatalf("Unable to read file (%s): %s", filepath, err.Error())
			}
			if len(f) == 0 {
				log.Fatalf("Emtpy file (%s)", filepath)
			}
			ns = strings.Split(strings.Trim(string(f), "\n"), "\n")
		} else {
			ns = strings.Split(*servers_string, ",")
		}
		for i, s := range ns {
			if !strings.Contains(s, ":") {
				ns[i] = strings.TrimSpace(s) + ":53"
			}
		}
		gc.NameServers = ns
		gc.NameServersSpecified = true
	}
	if *nanoSeconds {
		gc.TimeFormat = time.RFC3339Nano
	} else {
		gc.TimeFormat = time.RFC3339
	}
	if gc.GoMaxProcs < 0 {
		log.Fatal("Invalid argument for --go-processes. Must be >1.")
	}
	if gc.GoMaxProcs != 0 {
		runtime.GOMAXPROCS(gc.GoMaxProcs)
	}
	if gc.UDPOnly && gc.TCPOnly {
		log.Fatal("TCP Only and UDP Only are conflicting")
	}
	// Output Groups are defined by a base + any additional fields that the user wants
	groups := strings.Split(gc.IncludeInOutput, ",")
	if gc.ResultVerbosity != "short" && gc.ResultVerbosity != "normal" && gc.ResultVerbosity != "long" && gc.ResultVerbosity != "trace" {
		log.Fatal("Invalid result verbosity. Options: short, normal, long, trace")
	}

	gc.OutputGroups = append(gc.OutputGroups, gc.ResultVerbosity)
	gc.OutputGroups = append(gc.OutputGroups, groups...)

	if len(flags.Args()) > 0 {
		stat, _ := os.Stdin.Stat()
		// If stdin is piped from the terminal, and we havent specified a file, and if we have unparsed args
		// use them for a dig like reslution
		if (stat.Mode()&os.ModeCharDevice) != 0 && gc.InputFilePath == "-" && len(flags.Args()) == 1 {
			gc.PassedName = flags.Args()[0]
		} else {
			log.Fatal("Unused command line flags: ", flags.Args())
		}
	}

	// some modules require multiple passes over a file (this is really just the case for zone files)
	if !factory.AllowStdIn() && gc.InputFilePath == "-" {
		log.Fatal("Specified module does not allow reading from stdin")
	}

	// allow the factory to initialize itself
	if err := factory.Initialize(&gc); err != nil {
		log.Fatal("Factory was unable to initialize:", err.Error())
	}
	// run it.
	if err := zdns.DoLookups(&factory, &gc); err != nil {
		log.Fatal("Unable to run lookups:", err.Error())
	}
	// allow the factory to initialize itself
	if err := factory.Finalize(); err != nil {
		log.Fatal("Factory was unable to finalize:", err.Error())
	}
}
<|endoftext|>


<commit_before>package dns

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
)

// Parse the $GENERATE statement as used in BIND9 zones.
// See http://www.zytrax.com/books/dns/ch8/generate.html for instance.
// We are called after '$GENERATE '. After which we expect:
// * the range (12-24/2)
// * lhs (ownername)
// * [[ttl][class]]
// * type
// * rhs (rdata)
func generate(l lex, c chan lex, t chan Token, o string) string {
	step := 1
	if i := strings.IndexAny(l.token, "/"); i != -1 {
		if i+1 == len(l.token) {
			return "bad step in $GENERATE range"
		}
		if s, e := strconv.Atoi(l.token[i+1:]); e != nil {
			return "bad step in $GENERATE range"
		} else {
			if s < 0 {
				return "bad step in $GENERATE range"
			}
			step = s
		}
		l.token = l.token[:i]
	}
	sx := strings.SplitN(l.token, "-", 2)
	if len(sx) != 2 {
		return "bad start-stop in $GENERATE range"
	}
	start, err := strconv.Atoi(sx[0])
	if err != nil {
		return "bad start in $GENERATE range"
	}
	end, err := strconv.Atoi(sx[1])
	if err != nil {
		return "bad stop in $GENERATE range"
	}
	if end < 0 || start < 0 || end <= start {
		return "bad range in $GENERATE range"
	}

	<-c // _BLANK
	// Create a complete new string, which we then parse again.
	s := ""
BuildRR:
	l = <-c
	if l.value != _NEWLINE && l.value != _EOF {
		s += l.token
		goto BuildRR
	}
	for i := start; i <= end; i += step {
		var (
			escape bool
			dom string
			mod string
			offset int
			err error
		)

		for j := 0; j < len(s); j++ { // No 'range' because we need to jump around
			switch s[j] {
			case '\\':
				if escape {
					dom += "\\"
					escape = false
					continue
				}
				escape = true
			case '$':
				mod = "%d"
				offset = 0
				if escape {
					dom += "$"
					escape = false
					continue
				}
				escape = false
				if j+1 >= len(s) { // End of the string
					dom += fmt.Sprintf(mod, i+offset)
					continue
				} else {
					if s[j+1] == '$' {
						dom += "$"
						j++
						continue
					}
				}
				// Search for { and }
				if s[j+1] == '{' { // Modifier block
					sep := strings.Index(s[j+2:], "}")
					if sep == -1 {
						return "bad modifier in $GENERATE"
					}
					//println("checking", s[j+2:j+2+sep])
					mod, offset, err = modToPrintf(s[j+2 : j+2+sep])
					if err != nil {
						return "bad modifier in $GENERATE"
					}
					j += 2 + sep // Jump to it
				}
				//println("mod", mod)
				dom += fmt.Sprintf(mod, i+offset)
			default:
				if escape { // Pretty useless here
					escape = false
					continue
				}
				dom += string(s[j])
			}
		}
		// Re-parse the RR and send it on the current channel t
		rx, err := NewRR("$ORIGIN " + o + "\n" + dom)
		if err != nil {
			return err.(*ParseError).err
		}
		t <- Token{RR: rx}
	}
	return ""
}

// Convert a $GENERATE modifier 0,0,d to something Printf can deal with.
func modToPrintf(s string) (string, int, error) {
	xs := strings.SplitN(s, ",", 3)
	if len(xs) != 3 {
		return "", 0, errors.New("fubar")
	}
	// xs[0] is offset, xs[1] is width, xs[2] is base
	if xs[2] != "o" && xs[2] != "d" && xs[2] != "x" && xs[2] != "X" {
		return "", 0, errors.New("fubar")
	}
	offset, err := strconv.Atoi(xs[0])
	if err != nil {
		return "", 0, err
	}
	width, err := strconv.Atoi(xs[1])
	if err != nil {
		return "", offset, err
	}
	printf := "%"
	switch {
	case width < 0:
		return "", offset, errors.New("fubar")
	case width == 0:
		printf += xs[1]
	default:
		printf += "0" + xs[1]
	}
	printf += xs[2]
	return printf, offset, nil
}
<commit_msg>more descriptive errors<commit_after>package dns

import (
	"fmt"
	"strconv"
	"strings"
)

// Parse the $GENERATE statement as used in BIND9 zones.
// See http://www.zytrax.com/books/dns/ch8/generate.html for instance.
// We are called after '$GENERATE '. After which we expect:
// * the range (12-24/2)
// * lhs (ownername)
// * [[ttl][class]]
// * type
// * rhs (rdata)
// But we are lazy here, only the range is parsed *all* occurences
// of $ after that are interpreted.
// Any error are returned as a string value, the empty string signals
// "no error".
func generate(l lex, c chan lex, t chan Token, o string) string {
	step := 1
	if i := strings.IndexAny(l.token, "/"); i != -1 {
		if i+1 == len(l.token) {
			return "bad step in $GENERATE range"
		}
		if s, e := strconv.Atoi(l.token[i+1:]); e != nil {
			return "bad step in $GENERATE range"
		} else {
			if s < 0 {
				return "bad step in $GENERATE range"
			}
			step = s
		}
		l.token = l.token[:i]
	}
	sx := strings.SplitN(l.token, "-", 2)
	if len(sx) != 2 {
		return "bad start-stop in $GENERATE range"
	}
	start, err := strconv.Atoi(sx[0])
	if err != nil {
		return "bad start in $GENERATE range"
	}
	end, err := strconv.Atoi(sx[1])
	if err != nil {
		return "bad stop in $GENERATE range"
	}
	if end < 0 || start < 0 || end <= start {
		return "bad range in $GENERATE range"
	}

	<-c // _BLANK
	// Create a complete new string, which we then parse again.
	s := ""
BuildRR:
	l = <-c
	if l.value != _NEWLINE && l.value != _EOF {
		s += l.token
		goto BuildRR
	}
	for i := start; i <= end; i += step {
		var (
			escape bool
			dom string
			mod string
			err string
			offset int
		)

		for j := 0; j < len(s); j++ { // No 'range' because we need to jump around
			switch s[j] {
			case '\\':
				if escape {
					dom += "\\"
					escape = false
					continue
				}
				escape = true
			case '$':
				mod = "%d"
				offset = 0
				if escape {
					dom += "$"
					escape = false
					continue
				}
				escape = false
				if j+1 >= len(s) { // End of the string
					dom += fmt.Sprintf(mod, i+offset)
					continue
				} else {
					if s[j+1] == '$' {
						dom += "$"
						j++
						continue
					}
				}
				// Search for { and }
				if s[j+1] == '{' { // Modifier block
					sep := strings.Index(s[j+2:], "}")
					if sep == -1 {
						return "bad modifier in $GENERATE"
					}
					//println("checking", s[j+2:j+2+sep])
					mod, offset, err = modToPrintf(s[j+2 : j+2+sep])
					if err != "" {
						return err
					}
					j += 2 + sep // Jump to it
				}
				//println("mod", mod)
				dom += fmt.Sprintf(mod, i+offset)
			default:
				if escape { // Pretty useless here
					escape = false
					continue
				}
				dom += string(s[j])
			}
		}
		// Re-parse the RR and send it on the current channel t
		rx, e := NewRR("$ORIGIN " + o + "\n" + dom)
		if e != nil {
			return e.(*ParseError).err
		}
		t <- Token{RR: rx}
	}
	return ""
}

// Convert a $GENERATE modifier 0,0,d to something Printf can deal with.
func modToPrintf(s string) (string, int, string) {
	xs := strings.SplitN(s, ",", 3)
	if len(xs) != 3 {
		return "", 0, "bad modifier in $GENERATE"
	}
	// xs[0] is offset, xs[1] is width, xs[2] is base
	if xs[2] != "o" && xs[2] != "d" && xs[2] != "x" && xs[2] != "X" {
		return "", 0, "bad base in $GENERATE"
	}
	offset, err := strconv.Atoi(xs[0])
	if err != nil {
		return "", 0, "bad offset in $GENERATE"
	}
	width, err := strconv.Atoi(xs[1])
	if err != nil {
		return "", offset, "bad width in $GENERATE"
	}
	printf := "%"
	switch {
	case width < 0:
		return "", offset, "bad width in $GENERATE"
	case width == 0:
		printf += xs[1]
	default:
		printf += "0" + xs[1]
	}
	printf += xs[2]
	return printf, offset, ""
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package state_test

import (
	"errors"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/backup/mock"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/fs/mock"
	"github.com/jacobsa/comeback/state"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"testing"
	"time"
)

func TestScoreMapSaver(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ScoreMapSaverTest struct {
	sourceMap  state.ScoreMap
	sinkMap    state.ScoreMap
	fileSystem mock_fs.MockFileSystem
	wrapped    mock_backup.MockFileSaver
	saver      backup.FileSaver

	path   string
	scores []blob.Score
	err    error
}

func init() { RegisterTestSuite(&ScoreMapSaverTest{}) }

func (t *ScoreMapSaverTest) SetUp(i *TestInfo) {
	t.sourceMap = state.NewScoreMap()
	t.sinkMap = state.NewScoreMap()
	t.fileSystem = mock_fs.NewMockFileSystem(i.MockController, "fileSystem")
	t.wrapped = mock_backup.NewMockFileSaver(i.MockController, "wrapped")

	t.saver = state.NewScoreMapFileSaver(
		t.sourceMap,
		t.sinkMap,
		t.fileSystem,
		t.wrapped,
	)
}

func (t *ScoreMapSaverTest) call() {
	t.scores, t.err = t.saver.Save(t.path)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ScoreMapSaverTest) CallsStat() {
	t.path = "taco"

	// File system
	ExpectCall(t.fileSystem, "Stat")("taco").
		WillOnce(oglemock.Return(fs.DirectoryEntry{}, errors.New("")))

	// Call
	t.call()
}

func (t *ScoreMapSaverTest) StatReturnsError() {
	// File system
	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(fs.DirectoryEntry{}, errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("Stat")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *ScoreMapSaverTest) ScoreMapContainsEntry() {
	t.path = "taco"

	expectedKey := state.ScoreMapKey{
		Path:        "taco",
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       time.Now(),
		Inode:       23,
		Size:        29,
	}

	// Source map
	expectedScores := []blob.Score{
		blob.ComputeScore([]byte("foo")),
		blob.ComputeScore([]byte("bar")),
	}

	t.sourceMap.Set(expectedKey, expectedScores)

	// File system
	entry := fs.DirectoryEntry{
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       expectedKey.MTime,
		Inode:       23,
		Size:        29,
	}

	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(entry, nil))

	// Call
	t.call()

	AssertEq(nil, t.err)
	ExpectThat(t.scores, DeepEquals(expectedScores))
	ExpectThat(t.sinkMap.Get(expectedKey), DeepEquals(expectedScores))
}

func (t *ScoreMapSaverTest) CallsWrapped() {
	t.path = "taco"

	// File system
	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(fs.DirectoryEntry{}, nil))

	// Wrapped
	ExpectCall(t.wrapped, "Save")("taco").
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *ScoreMapSaverTest) WrappedReturnsError() {
	expectedKey := state.ScoreMapKey{
		Path:        "taco",
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       time.Now(),
		Inode:       23,
		Size:        29,
	}

	// File system
	entry := fs.DirectoryEntry{
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       expectedKey.MTime,
		Inode:       23,
		Size:        29,
	}

	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(entry, nil))

	// Wrapped
	ExpectCall(t.wrapped, "Save")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.call()

	AssertThat(t.err, Error(Equals("taco")))
	ExpectEq(nil, t.sinkMap.Get(expectedKey))
}

func (t *ScoreMapSaverTest) WrappedReturnsScores() {
	t.path = "taco"
	expectedKey := state.ScoreMapKey{
		Path:        "taco",
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       time.Now(),
		Inode:       23,
		Size:        29,
	}

	// File system
	entry := fs.DirectoryEntry{
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       expectedKey.MTime,
		Inode:       23,
		Size:        29,
	}

	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(entry, nil))

	// Wrapped
	expectedScores := []blob.Score{
		blob.ComputeScore([]byte("foo")),
		blob.ComputeScore([]byte("bar")),
	}

	ExpectCall(t.wrapped, "Save")(Any()).
		WillOnce(oglemock.Return(expectedScores, nil))

	// Call
	t.call()

	AssertEq(nil, t.err)
	ExpectThat(t.sinkMap.Get(expectedKey), DeepEquals(expectedScores))
}
<commit_msg>Changed test package.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package state

import (
	"errors"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/backup/mock"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/fs/mock"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"testing"
	"time"
)

func TestScoreMapSaver(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ScoreMapSaverTest struct {
	sourceMap  ScoreMap
	sinkMap    ScoreMap
	fileSystem mock_fs.MockFileSystem
	wrapped    mock_backup.MockFileSaver
	saver      backup.FileSaver

	path   string
	scores []blob.Score
	err    error
}

func init() { RegisterTestSuite(&ScoreMapSaverTest{}) }

func (t *ScoreMapSaverTest) SetUp(i *TestInfo) {
	t.sourceMap = NewScoreMap()
	t.sinkMap = NewScoreMap()
	t.fileSystem = mock_fs.NewMockFileSystem(i.MockController, "fileSystem")
	t.wrapped = mock_backup.NewMockFileSaver(i.MockController, "wrapped")

	t.saver = NewScoreMapFileSaver(
		t.sourceMap,
		t.sinkMap,
		t.fileSystem,
		t.wrapped,
	)
}

func (t *ScoreMapSaverTest) call() {
	t.scores, t.err = t.saver.Save(t.path)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ScoreMapSaverTest) CallsStat() {
	t.path = "taco"

	// File system
	ExpectCall(t.fileSystem, "Stat")("taco").
		WillOnce(oglemock.Return(fs.DirectoryEntry{}, errors.New("")))

	// Call
	t.call()
}

func (t *ScoreMapSaverTest) StatReturnsError() {
	// File system
	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(fs.DirectoryEntry{}, errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("Stat")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *ScoreMapSaverTest) ScoreMapContainsEntry() {
	t.path = "taco"

	expectedKey := ScoreMapKey{
		Path:        "taco",
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       time.Now(),
		Inode:       23,
		Size:        29,
	}

	// Source map
	expectedScores := []blob.Score{
		blob.ComputeScore([]byte("foo")),
		blob.ComputeScore([]byte("bar")),
	}

	t.sourceMap.Set(expectedKey, expectedScores)

	// File system
	entry := fs.DirectoryEntry{
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       expectedKey.MTime,
		Inode:       23,
		Size:        29,
	}

	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(entry, nil))

	// Call
	t.call()

	AssertEq(nil, t.err)
	ExpectThat(t.scores, DeepEquals(expectedScores))
	ExpectThat(t.sinkMap.Get(expectedKey), DeepEquals(expectedScores))
}

func (t *ScoreMapSaverTest) CallsWrapped() {
	t.path = "taco"

	// File system
	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(fs.DirectoryEntry{}, nil))

	// Wrapped
	ExpectCall(t.wrapped, "Save")("taco").
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *ScoreMapSaverTest) WrappedReturnsError() {
	expectedKey := ScoreMapKey{
		Path:        "taco",
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       time.Now(),
		Inode:       23,
		Size:        29,
	}

	// File system
	entry := fs.DirectoryEntry{
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       expectedKey.MTime,
		Inode:       23,
		Size:        29,
	}

	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(entry, nil))

	// Wrapped
	ExpectCall(t.wrapped, "Save")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.call()

	AssertThat(t.err, Error(Equals("taco")))
	ExpectEq(nil, t.sinkMap.Get(expectedKey))
}

func (t *ScoreMapSaverTest) WrappedReturnsScores() {
	t.path = "taco"
	expectedKey := ScoreMapKey{
		Path:        "taco",
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       time.Now(),
		Inode:       23,
		Size:        29,
	}

	// File system
	entry := fs.DirectoryEntry{
		Permissions: 0644,
		Uid:         17,
		Gid:         19,
		MTime:       expectedKey.MTime,
		Inode:       23,
		Size:        29,
	}

	ExpectCall(t.fileSystem, "Stat")(Any()).
		WillOnce(oglemock.Return(entry, nil))

	// Wrapped
	expectedScores := []blob.Score{
		blob.ComputeScore([]byte("foo")),
		blob.ComputeScore([]byte("bar")),
	}

	ExpectCall(t.wrapped, "Save")(Any()).
		WillOnce(oglemock.Return(expectedScores, nil))

	// Call
	t.call()

	AssertEq(nil, t.err)
	ExpectThat(t.sinkMap.Get(expectedKey), DeepEquals(expectedScores))
}
<|endoftext|>


<commit_before>package zlib

import (
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"regexp"
	"time"

	"ztools/ztls"
)

var smtpEndRegex = regexp.MustCompile(`(?:\r\n)|^[0-9]{3} .+\r\n$`)
var pop3EndRegex = regexp.MustCompile(`(?:\r\n\.\r\n$)|(?:\r\n$)`)
var imapStatusEndRegex = regexp.MustCompile(`\r\n$`)

const (
	SMTP_COMMAND = "STARTTLS\r\n"
	POP3_COMMAND = "STLS\r\n"
	IMAP_COMMAND = "a001 STARTTLS\r\n"
)

// Implements the net.Conn interface
type Conn struct {
	// Underlying network connection
	conn    net.Conn
	tlsConn *ztls.Conn
	isTls   bool

	// Max TLS version
	maxTlsVersion uint16

	// Keep track of state / network operations
	operations []ConnectionEvent

	// Cache the deadlines so we can reapply after TLS handshake
	readDeadline  time.Time
	writeDeadline time.Time

	caPool  *x509.CertPool
	onlyCBC bool

	domain string
}

func (c *Conn) getUnderlyingConn() net.Conn {
	if c.isTls {
		return c.tlsConn
	}
	return c.conn
}

func (c *Conn) SetCBCOnly() {
	c.onlyCBC = true
}

func (c *Conn) SetCAPool(pool *x509.CertPool) {
	c.caPool = pool
}

func (c *Conn) SetDomain(domain string) {
	c.domain = domain
}

// Layer in the regular conn methods
func (c *Conn) LocalAddr() net.Addr {
	return c.getUnderlyingConn().LocalAddr()
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.getUnderlyingConn().RemoteAddr()
}

func (c *Conn) SetDeadline(t time.Time) error {
	c.readDeadline = t
	c.writeDeadline = t
	return c.getUnderlyingConn().SetDeadline(t)
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	c.readDeadline = t
	return c.getUnderlyingConn().SetReadDeadline(t)
}

func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.writeDeadline = t
	return c.getUnderlyingConn().SetWriteDeadline(t)
}

// Delegate here, but record all the things
func (c *Conn) Write(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Write(b)
	w := WriteEvent{Sent: b}
	event := ConnectionEvent{
		Data:  &w,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) Read(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Read(b)
	r := ReadEvent{
		Response: b[0:n],
	}
	c.appendEvent(&r, err)
	return n, err
}

func (c *Conn) Close() error {
	return c.getUnderlyingConn().Close()
}

// Extra method - Do a TLS Handshake and record progress
func (c *Conn) TLSHandshake() error {
	if c.isTls {
		return fmt.Errorf(
			"Attempted repeat handshake with remote host %s",
			c.RemoteAddr().String())
	}
	tlsConfig := new(ztls.Config)
	tlsConfig.InsecureSkipVerify = true
	tlsConfig.MinVersion = ztls.VersionSSL30
	tlsConfig.MaxVersion = c.maxTlsVersion
	tlsConfig.RootCAs = c.caPool
	if c.domain != "" {
		tlsConfig.ServerName = c.domain
	}
	if c.onlyCBC {
		tlsConfig.CipherSuites = ztls.CBCSuiteIDList
	}
	c.tlsConn = ztls.Client(c.conn, tlsConfig)
	c.tlsConn.SetReadDeadline(c.readDeadline)
	c.tlsConn.SetWriteDeadline(c.writeDeadline)
	c.isTls = true
	err := c.tlsConn.Handshake()
	hl := c.tlsConn.GetHandshakeLog()
	ts := TLSHandshakeEvent{handshakeLog: hl}
	event := ConnectionEvent{
		Data:  &ts,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return err
}

func (c *Conn) sendStartTLSCommand(command string) error {
	// Don't doublehandshake
	if c.isTls {
		return fmt.Errorf(
			"Attempt STARTTLS after TLS handshake with remote host %s",
			c.RemoteAddr().String())
	}
	// Send the STARTTLS message
	starttls := []byte(command)
	_, err := c.conn.Write(starttls)
	return err
}

// Do a STARTTLS handshake
func (c *Conn) SMTPStartTLSHandshake() error {
	// Make the state
	ss := StartTLSEvent{Command: SMTP_COMMAND}

	// Send the command
	if err := c.sendStartTLSCommand(SMTP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}
	// Read the response on a successful send
	buf := make([]byte, 256)
	n, err := c.readSmtpResponse(buf)
	ss.Response = buf[0:n]

	// Record everything no matter the result
	c.appendEvent(&ss, err)

	// Stop if we failed already
	if err != nil {
		return err
	}
	// Successful so far, attempt to do the actual handshake
	return c.TLSHandshake()
}

func (c *Conn) POP3StartTLSHandshake() error {
	ss := StartTLSEvent{Command: POP3_COMMAND}
	if err := c.sendStartTLSCommand(POP3_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readPop3Response(buf)
	ss.Response = buf[0:n]
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) IMAPStartTLSHandshake() error {
	ss := StartTLSEvent{Command: IMAP_COMMAND}
	if err := c.sendStartTLSCommand(IMAP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readImapStatusResponse(buf)
	ss.Response = buf[0:n]
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) readUntilRegex(res []byte, expr *regexp.Regexp) (int, error) {
	buf := res[0:]
	length := 0
	for finished := false; !finished; {
		n, err := c.getUnderlyingConn().Read(buf)
		length += n
		if err != nil {
			return length, err
		}
		if expr.Match(res[0:length]) {
			finished = true
		}
		if length == len(res) {
			return length, errors.New("Not enough buffer space")
		}
		buf = res[length:]
	}
	return length, nil
}

func (c *Conn) readSmtpResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, smtpEndRegex)
}

func (c *Conn) SMTPBanner(b []byte) (int, error) {
	n, err := c.readSmtpResponse(b)
	mb := MailBannerEvent{}
	mb.Banner = string(b[0:n])
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) EHLO(domain string) error {
	cmd := []byte("EHLO " + domain + "\r\n")
	ee := EHLOEvent{}
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(&ee, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	ee.Response = buf[0:n]
	c.appendEvent(&ee, err)
	return err
}

func (c *Conn) SMTPHelp() error {
	cmd := []byte("HELP\r\n")
	h := new(SMTPHelpEvent)
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(h, err)
		return err
	}
	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	h.Response = buf[0:n]
	c.appendEvent(h, err)
	return err
}

func (c *Conn) readPop3Response(res []byte) (int, error) {
	return c.readUntilRegex(res, pop3EndRegex)
}

func (c *Conn) POP3Banner(b []byte) (int, error) {
	n, err := c.readPop3Response(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) readImapStatusResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, imapStatusEndRegex)
}

func (c *Conn) IMAPBanner(b []byte) (int, error) {
	n, err := c.readImapStatusResponse(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) CheckHeartbleed(b []byte) (int, error) {
	if !c.isTls {
		return 0, fmt.Errorf(
			"Must perform TLS handshake before sending Heartbleed probe to %s",
			c.RemoteAddr().String())
	}
	n, err := c.tlsConn.CheckHeartbleed(b)
	hb := c.tlsConn.GetHeartbleedLog()
	if err == ztls.HeartbleedError {
		err = nil
	}
	event := ConnectionEvent{
		Data:  &HeartbleedEvent{heartbleedLog: hb},
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) SendModbusEcho(b []byte) (int, error) {
	// TODO MODBUS
	return 0, nil
}

func (c *Conn) States() []ConnectionEvent {
	return c.operations
}

func (c *Conn) appendEvent(data EventData, err error) {
	event := ConnectionEvent{
		Data:  data,
		Error: err,
	}
	c.operations = append(c.operations, event)
}
<commit_msg>Clarify modbus TODO<commit_after>package zlib

import (
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"regexp"
	"time"

	"ztools/ztls"
)

var smtpEndRegex = regexp.MustCompile(`(?:\r\n)|^[0-9]{3} .+\r\n$`)
var pop3EndRegex = regexp.MustCompile(`(?:\r\n\.\r\n$)|(?:\r\n$)`)
var imapStatusEndRegex = regexp.MustCompile(`\r\n$`)

const (
	SMTP_COMMAND = "STARTTLS\r\n"
	POP3_COMMAND = "STLS\r\n"
	IMAP_COMMAND = "a001 STARTTLS\r\n"
)

// Implements the net.Conn interface
type Conn struct {
	// Underlying network connection
	conn    net.Conn
	tlsConn *ztls.Conn
	isTls   bool

	// Max TLS version
	maxTlsVersion uint16

	// Keep track of state / network operations
	operations []ConnectionEvent

	// Cache the deadlines so we can reapply after TLS handshake
	readDeadline  time.Time
	writeDeadline time.Time

	caPool  *x509.CertPool
	onlyCBC bool

	domain string
}

func (c *Conn) getUnderlyingConn() net.Conn {
	if c.isTls {
		return c.tlsConn
	}
	return c.conn
}

func (c *Conn) SetCBCOnly() {
	c.onlyCBC = true
}

func (c *Conn) SetCAPool(pool *x509.CertPool) {
	c.caPool = pool
}

func (c *Conn) SetDomain(domain string) {
	c.domain = domain
}

// Layer in the regular conn methods
func (c *Conn) LocalAddr() net.Addr {
	return c.getUnderlyingConn().LocalAddr()
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.getUnderlyingConn().RemoteAddr()
}

func (c *Conn) SetDeadline(t time.Time) error {
	c.readDeadline = t
	c.writeDeadline = t
	return c.getUnderlyingConn().SetDeadline(t)
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	c.readDeadline = t
	return c.getUnderlyingConn().SetReadDeadline(t)
}

func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.writeDeadline = t
	return c.getUnderlyingConn().SetWriteDeadline(t)
}

// Delegate here, but record all the things
func (c *Conn) Write(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Write(b)
	w := WriteEvent{Sent: b}
	event := ConnectionEvent{
		Data:  &w,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) Read(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Read(b)
	r := ReadEvent{
		Response: b[0:n],
	}
	c.appendEvent(&r, err)
	return n, err
}

func (c *Conn) Close() error {
	return c.getUnderlyingConn().Close()
}

// Extra method - Do a TLS Handshake and record progress
func (c *Conn) TLSHandshake() error {
	if c.isTls {
		return fmt.Errorf(
			"Attempted repeat handshake with remote host %s",
			c.RemoteAddr().String())
	}
	tlsConfig := new(ztls.Config)
	tlsConfig.InsecureSkipVerify = true
	tlsConfig.MinVersion = ztls.VersionSSL30
	tlsConfig.MaxVersion = c.maxTlsVersion
	tlsConfig.RootCAs = c.caPool
	if c.domain != "" {
		tlsConfig.ServerName = c.domain
	}
	if c.onlyCBC {
		tlsConfig.CipherSuites = ztls.CBCSuiteIDList
	}
	c.tlsConn = ztls.Client(c.conn, tlsConfig)
	c.tlsConn.SetReadDeadline(c.readDeadline)
	c.tlsConn.SetWriteDeadline(c.writeDeadline)
	c.isTls = true
	err := c.tlsConn.Handshake()
	hl := c.tlsConn.GetHandshakeLog()
	ts := TLSHandshakeEvent{handshakeLog: hl}
	event := ConnectionEvent{
		Data:  &ts,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return err
}

func (c *Conn) sendStartTLSCommand(command string) error {
	// Don't doublehandshake
	if c.isTls {
		return fmt.Errorf(
			"Attempt STARTTLS after TLS handshake with remote host %s",
			c.RemoteAddr().String())
	}
	// Send the STARTTLS message
	starttls := []byte(command)
	_, err := c.conn.Write(starttls)
	return err
}

// Do a STARTTLS handshake
func (c *Conn) SMTPStartTLSHandshake() error {
	// Make the state
	ss := StartTLSEvent{Command: SMTP_COMMAND}

	// Send the command
	if err := c.sendStartTLSCommand(SMTP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}
	// Read the response on a successful send
	buf := make([]byte, 256)
	n, err := c.readSmtpResponse(buf)
	ss.Response = buf[0:n]

	// Record everything no matter the result
	c.appendEvent(&ss, err)

	// Stop if we failed already
	if err != nil {
		return err
	}
	// Successful so far, attempt to do the actual handshake
	return c.TLSHandshake()
}

func (c *Conn) POP3StartTLSHandshake() error {
	ss := StartTLSEvent{Command: POP3_COMMAND}
	if err := c.sendStartTLSCommand(POP3_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readPop3Response(buf)
	ss.Response = buf[0:n]
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) IMAPStartTLSHandshake() error {
	ss := StartTLSEvent{Command: IMAP_COMMAND}
	if err := c.sendStartTLSCommand(IMAP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readImapStatusResponse(buf)
	ss.Response = buf[0:n]
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) readUntilRegex(res []byte, expr *regexp.Regexp) (int, error) {
	buf := res[0:]
	length := 0
	for finished := false; !finished; {
		n, err := c.getUnderlyingConn().Read(buf)
		length += n
		if err != nil {
			return length, err
		}
		if expr.Match(res[0:length]) {
			finished = true
		}
		if length == len(res) {
			return length, errors.New("Not enough buffer space")
		}
		buf = res[length:]
	}
	return length, nil
}

func (c *Conn) readSmtpResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, smtpEndRegex)
}

func (c *Conn) SMTPBanner(b []byte) (int, error) {
	n, err := c.readSmtpResponse(b)
	mb := MailBannerEvent{}
	mb.Banner = string(b[0:n])
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) EHLO(domain string) error {
	cmd := []byte("EHLO " + domain + "\r\n")
	ee := EHLOEvent{}
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(&ee, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	ee.Response = buf[0:n]
	c.appendEvent(&ee, err)
	return err
}

func (c *Conn) SMTPHelp() error {
	cmd := []byte("HELP\r\n")
	h := new(SMTPHelpEvent)
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(h, err)
		return err
	}
	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	h.Response = buf[0:n]
	c.appendEvent(h, err)
	return err
}

func (c *Conn) readPop3Response(res []byte) (int, error) {
	return c.readUntilRegex(res, pop3EndRegex)
}

func (c *Conn) POP3Banner(b []byte) (int, error) {
	n, err := c.readPop3Response(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) readImapStatusResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, imapStatusEndRegex)
}

func (c *Conn) IMAPBanner(b []byte) (int, error) {
	n, err := c.readImapStatusResponse(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) CheckHeartbleed(b []byte) (int, error) {
	if !c.isTls {
		return 0, fmt.Errorf(
			"Must perform TLS handshake before sending Heartbleed probe to %s",
			c.RemoteAddr().String())
	}
	n, err := c.tlsConn.CheckHeartbleed(b)
	hb := c.tlsConn.GetHeartbleedLog()
	if err == ztls.HeartbleedError {
		err = nil
	}
	event := ConnectionEvent{
		Data:  &HeartbleedEvent{heartbleedLog: hb},
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) SendModbusEcho(b []byte) (int, error) {
	// TODO MODBUS
	// Build a request

	// Send it out

	// Read the response

	// make sure the whole thing gets appended to the operation log
	// e.g. c.appendEvent(modbusEvent, modbusError)
	return 0, nil
}

func (c *Conn) States() []ConnectionEvent {
	return c.operations
}

func (c *Conn) appendEvent(data EventData, err error) {
	event := ConnectionEvent{
		Data:  data,
		Error: err,
	}
	c.operations = append(c.operations, event)
}
<|endoftext|>


<commit_before>package common

import "fmt"

const (
	MajorVersion      = 0
	MinorVersion      = 9
	MicroVersion      = 0
	AdditionalVersion = ""
)

func VersionString() string {
	return fmt.Sprintf("%d.%d.%d%s", MajorVersion, MinorVersion, MicroVersion, AdditionalVersion)
}
<commit_msg>bump master version to 0.10.0dev<commit_after>package common

import "fmt"

const (
	MajorVersion      = 0
	MinorVersion      = 10
	MicroVersion      = 0
	AdditionalVersion = "dev"
)

func VersionString() string {
	return fmt.Sprintf("%d.%d.%d%s", MajorVersion, MinorVersion, MicroVersion, AdditionalVersion)
}
<|endoftext|>


<commit_before>package common

import (
	"bytes"
    "encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
)

type Docker struct {
	h	       *httpAPI			// contains ip string
	Containers []Container
	Images     []Image
	Updated    time.Time
}

type Container struct {
	Id              string
	Image           string
	NetworkSettings struct {
		PortMapping struct {
			Tcp map[string]string
		}
	}
}

type Image struct {
	Id         string
	Repository string
}

type httpAPI struct {
	ip string
}

// function to initialize new httpAPI struct
func NewHttpClient(ip string) (h *httpAPI) {
	h = &httpAPI{ip}
	return
}

// function to initialize new Docker struct
func NewDocker(ip string) (D *Docker) {
	D = &Docker{NewHttpClient(ip)}
	return
}

// Post function to clean up http.Post calls in code, method for HttpAPI struct
func (h *httpAPI) Post(url string, content string, b io.Reader) (resp *http.Response, err error) {
	c := MakeHttpClient()

	resp, err = c.Post("http://"+h.ip+":4243/"+url,
		content, b)

	return
}

// Post with a dictionary of header values
func (h *httpAPI) PostHeader(url string, content string, b io.Reader, header http.Header) (resp *http.Response, err error) {
	c := MakeHttpClient()

	req, err := http.NewRequest("POST",
		"http://"+h.ip+":4243/"+url, b)
	if err != nil {
		log.Fatal(err)
	}

    req.Header = header
    req.Header.Set("Content-Type", content)

	resp, err = c.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	return
}

// Get function to clean up http.Get calls in code, method for HttpAPI struct
func (h *httpAPI) Get(url string) (resp *http.Response, err error) {
	c := MakeHttpClient()

	resp, err = c.Get("http://" + h.ip + ":4243/" + url)

	return
}

// Delete function to clean up http.NewRequest("DELETE"...) call, method for HttpAPI struct
func (h *httpAPI) Delete(url string) (resp *http.Response, err error) {
	c := MakeHttpClient()

	req, err := http.NewRequest("DELETE",
		"http://"+h.ip+":4243/"+url, nil)
	if err != nil {
		log.Fatal(err)
	}
	resp, err = c.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	return
}

func (D *Docker) GetIP() string {
    return D.h.ip
}

// function to periodically update information in Docker struct
func (D *Docker) Update() {
	for ; ; time.Sleep(60 * time.Second) {
		c, err := D.ListContainers()
		if err != nil {
			log.Print(err)
			continue
		}

		img, err := D.ListImages()
		if err != nil {
			log.Print(err)
			continue
		}

		D = Docker{_,c, img, time.Now()}
		update <- D
	}
}

// InspectContainer takes a container, and returns its port and its info
func (D *Docker) InspectContainer(id string) (err error) {
	C := &Container{}

	resp, err := D.h.Get("/containers/" + id + "/json")
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return errors.New("Inspect Container Status is not 200")
	}

	rall, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	json.Unmarshal(rall, C)

	log.Print("Container inspected")
	return nil
}

// runImage takes a docker image to run, and makes sure it is running
func (D *Docker) RunImage(imagename string, hint bool) (id string, err error) {

	e := make([]string, 1)
	e[0] = "HOST=" + D.h.ip

	c := make(map[string]interface{})
	c["Image"] = imagename
	if hint {
		c["Env"] = e
	}

	ba, err := json.Marshal(c)
	if err != nil {
		return "", err
	}
	b := bytes.NewBuffer(ba)

	resp, err := h.Post("containers/create", "application/json", b)

	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 201 {
		msg := fmt.Sprintf("Create Container Response status is %d", resp.StatusCode)
		return "", errors.New(msg)
	}

	s, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	a := &Container{}
	json.Unmarshal(s, a)
	id = a.Id

	log.Printf("Container created id:%s", id)

	b = bytes.NewBuffer([]byte("{}"))
	resp, err = D.h.Post("containers/"+id+"/start", "application/json", b)

	defer resp.Body.Close()

	LogReader(resp.Body)

	if resp.StatusCode != 204 {
		msg := fmt.Sprintf("Start Container Response status is %d", resp.StatusCode)
		return "", errors.New(msg)
	}

	log.Printf("Container running")
	return id, nil

}

// StopContainer stops a container
func (D *Docker) StopContainer(container string) (err error) {
	log.Print("Stopping container ", container)

	resp, err := D.h.Post("containers/"+container+"/stop?t=1", "application/json", b)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	LogReader(resp.Body)
	return nil
}

// StopImage takes a image to stop and stops it
func (D *Docker) StopImage(imagename string) {
	log.Print("Stopping image ", imagename)

	running, id, err := D.ImageRunning(imagename)
	if err != nil {
		log.Fatal(err)
	}

	if running {
		err = D.StopContainer(id)
		if err != nil {
			log.Fatal(err)
		}
		D.DeleteContainer(id)
	}
}

func (D *Docker) DeleteContainer(id string) (err error) {
	log.Print("deleting container ", id)

	resp, err := D.h.Delete("containers/" + id)
	if err != nil {
		return err
	}

	defer resp.Body.Close()
	LogReader(resp.Body)
	return
}

// TagImage tags an already existing image in the repository
func (D *Docker) TagImage(name string, tag string) (err error) {
	b := bytes.NewBuffer(nil)

	resp, err := D.h.Post("images/"+name+"/tag?repo="+tag+"&force=1", "application/json", b)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 201 {
		log.Print(resp)
		return errors.New("Code not 201")
	}

	return nil

}

// PushImage pushes an image to a docker index
func (D *Docker) PushImage(w io.Writer, name string) (err error) {
	b := bytes.NewBuffer(nil)

    auth := base64.StdEncoding.EncodeToString([]byte("{}"))
    header := make(http.Header)
    header.Set("X-Registry-Auth", auth)

    url := "images/"+name+"/push"

	resp, err := D.h.PostHeader(url,
		"application/json", b, header)
	if err != nil {
		return err
	}

	if resp.StatusCode != 200 {

		_, err = io.Copy(w, resp.Body)
        if err != nil {
            io.WriteString(w, err.Error())
        }

        io.WriteString(w, resp.Header.Get("Location")+"\n")
        io.WriteString(w, url + "\n")

		s := fmt.Sprintf("Push image status code is not 200 it is %d",
			resp.StatusCode)
		return errors.New(s)
	}

	defer resp.Body.Close()
	_, err = io.Copy(w, resp.Body)
	if err != nil {
		return err
	}

	return nil
}

// buildImage takes a tarfile, and builds it
func (D *Docker) BuildImage(fd io.Reader, name string) (err error) {

	v := fmt.Sprintf("%d", time.Now().Unix())
	resp, err := D.h.Post("build?t="+name+"%3A"+v,
		"application/tar", fd)

	if err != nil {
		return err
	}

	defer resp.Body.Close()

	LogReader(resp.Body)

	return D.TagImage(name+"%3A"+v, name)
}

// loadImage pulls a specified image into a docker instance
func (D *Docker) LoadImage(imagename string) (err error) {
	b := bytes.NewBuffer(nil)
	resp, err := D.h.Post("images/create?fromImage="+imagename,
		"text", b)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	LogReader(resp.Body)
	if resp.StatusCode != 200 {
		log.Printf("create status code is not 200 %s", resp.StatusCode)
	}

	log.Printf("Image fetched %s", imagename)
	return nil
}

// ListContainers gives the state for a specific docker container
func (D *Docker) ListContainers() (c []Container, err error) {
	resp, err := D.h.Get("containers/json")

	if err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return c, errors.New("Response code is not 200")
	}

	message, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	containers := &c
	err = json.Unmarshal(message, containers)
	return
}

// ListImages gives the state of the images for a specific Docker container
func (D *Docker) ListImages() (img []Image, err error) {
	resp, err := D.h.Get("images/json")

		if err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return img, errors.New("Response code is not 200")
	}

	message, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	images := &img
	err = json.Unmarshal(message, images)
	return
}
// ImageRunning states whether an image is running on a docker instance
func (D *Docker) ImageRunning(imagename string) (running bool, id string, err error) {

	containers, err := D.ListContainers()
	if err != nil {
		return false, "", err
	}

	for _, v := range containers {
		if strings.SplitN(v.Image, ":", 2)[0] == imagename {
			return true, v.Id, nil
		}
	}
	return
}
<commit_msg>common/docker: fixed compiler errors<commit_after>package common

import (
	"bytes"
    "encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
)

type Docker struct {
	h	       *httpAPI			// contains ip string
	Containers []Container
	Images     []Image
	Updated    time.Time
}

type Container struct {
	Id              string
	Image           string
	NetworkSettings struct {
		PortMapping struct {
			Tcp map[string]string
		}
	}
}

type Image struct {
	Id         string
	Repository string
}

type httpAPI struct {
	ip string
}

// function to initialize new httpAPI struct
func NewHttpClient(ip string) (h *httpAPI) {
	h = &httpAPI{ip}
	return
}

// function to initialize new Docker struct
func NewDocker(ip string) (D *Docker) {
	D = &Docker{}
	D.h = NewHttpClient(ip)
	return
}

// Post function to clean up http.Post calls in code, method for HttpAPI struct
func (h *httpAPI) Post(url string, content string, b io.Reader) (resp *http.Response, err error) {
	c := MakeHttpClient()

	resp, err = c.Post("http://"+h.ip+":4243/"+url,
		content, b)

	return
}

// Post with a dictionary of header values
func (h *httpAPI) PostHeader(url string, content string, b io.Reader, header http.Header) (resp *http.Response, err error) {
	c := MakeHttpClient()

	req, err := http.NewRequest("POST",
		"http://"+h.ip+":4243/"+url, b)
	if err != nil {
		log.Fatal(err)
	}

    req.Header = header
    req.Header.Set("Content-Type", content)

	resp, err = c.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	return
}

// Get function to clean up http.Get calls in code, method for HttpAPI struct
func (h *httpAPI) Get(url string) (resp *http.Response, err error) {
	c := MakeHttpClient()

	resp, err = c.Get("http://" + h.ip + ":4243/" + url)

	return
}

// Delete function to clean up http.NewRequest("DELETE"...) call, method for HttpAPI struct
func (h *httpAPI) Delete(url string) (resp *http.Response, err error) {
	c := MakeHttpClient()

	req, err := http.NewRequest("DELETE",
		"http://"+h.ip+":4243/"+url, nil)
	if err != nil {
		log.Fatal(err)
	}
	resp, err = c.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	return
}

func (D *Docker) GetIP() string {
    return D.h.ip
}

// function to periodically update information in Docker struct
func (D *Docker) Update() {
	for ; ; time.Sleep(60 * time.Second) {
		c, err := D.ListContainers()
		if err != nil {
			log.Print(err)
			continue
		}

		img, err := D.ListImages()
		if err != nil {
			log.Print(err)
			continue
		}

		D.Containers = c
		D.Images = img
		D.Updated = time.Now()
	}
}

// InspectContainer takes a container, and returns its port and its info
func (D *Docker) InspectContainer(id string) (err error) {
	C := &Container{}

	resp, err := D.h.Get("/containers/" + id + "/json")
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return errors.New("Inspect Container Status is not 200")
	}

	rall, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	json.Unmarshal(rall, C)

	log.Print("Container inspected")
	return nil
}

// runImage takes a docker image to run, and makes sure it is running
func (D *Docker) RunImage(imagename string, hint bool) (id string, err error) {

	e := make([]string, 1)
	e[0] = "HOST=" + D.h.ip

	c := make(map[string]interface{})
	c["Image"] = imagename
	if hint {
		c["Env"] = e
	}

	ba, err := json.Marshal(c)
	if err != nil {
		return "", err
	}
	b := bytes.NewBuffer(ba)

	resp, err := D.h.Post("containers/create", "application/json", b)

	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 201 {
		msg := fmt.Sprintf("Create Container Response status is %d", resp.StatusCode)
		return "", errors.New(msg)
	}

	s, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	a := &Container{}
	json.Unmarshal(s, a)
	id = a.Id

	log.Printf("Container created id:%s", id)

	b = bytes.NewBuffer([]byte("{}"))
	resp, err = D.h.Post("containers/"+id+"/start", "application/json", b)

	defer resp.Body.Close()

	LogReader(resp.Body)

	if resp.StatusCode != 204 {
		msg := fmt.Sprintf("Start Container Response status is %d", resp.StatusCode)
		return "", errors.New(msg)
	}

	log.Printf("Container running")
	return id, nil

}

// StopContainer stops a container
func (D *Docker) StopContainer(container string) (err error) {
	log.Print("Stopping container ", container)
	b := bytes.NewBuffer(nil)

	resp, err := D.h.Post("containers/"+container+"/stop?t=1", "application/json", b)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	LogReader(resp.Body)
	return nil
}

// StopImage takes a image to stop and stops it
func (D *Docker) StopImage(imagename string) {
	log.Print("Stopping image ", imagename)

	running, id, err := D.ImageRunning(imagename)
	if err != nil {
		log.Fatal(err)
	}

	if running {
		err = D.StopContainer(id)
		if err != nil {
			log.Fatal(err)
		}
		D.DeleteContainer(id)
	}
}

func (D *Docker) DeleteContainer(id string) (err error) {
	log.Print("deleting container ", id)

	resp, err := D.h.Delete("containers/" + id)
	if err != nil {
		return err
	}

	defer resp.Body.Close()
	LogReader(resp.Body)
	return
}

// TagImage tags an already existing image in the repository
func (D *Docker) TagImage(name string, tag string) (err error) {
	b := bytes.NewBuffer(nil)

	resp, err := D.h.Post("images/"+name+"/tag?repo="+tag+"&force=1", "application/json", b)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 201 {
		log.Print(resp)
		return errors.New("Code not 201")
	}

	return nil

}

// PushImage pushes an image to a docker index
func (D *Docker) PushImage(w io.Writer, name string) (err error) {
	b := bytes.NewBuffer(nil)

    auth := base64.StdEncoding.EncodeToString([]byte("{}"))
    header := make(http.Header)
    header.Set("X-Registry-Auth", auth)

    url := "images/"+name+"/push"

	resp, err := D.h.PostHeader(url,
		"application/json", b, header)
	if err != nil {
		return err
	}

	if resp.StatusCode != 200 {

		_, err = io.Copy(w, resp.Body)
        if err != nil {
            io.WriteString(w, err.Error())
        }

        io.WriteString(w, resp.Header.Get("Location")+"\n")
        io.WriteString(w, url + "\n")

		s := fmt.Sprintf("Push image status code is not 200 it is %d",
			resp.StatusCode)
		return errors.New(s)
	}

	defer resp.Body.Close()
	_, err = io.Copy(w, resp.Body)
	if err != nil {
		return err
	}

	return nil
}

// buildImage takes a tarfile, and builds it
func (D *Docker) BuildImage(fd io.Reader, name string) (err error) {

	v := fmt.Sprintf("%d", time.Now().Unix())
	resp, err := D.h.Post("build?t="+name+"%3A"+v,
		"application/tar", fd)

	if err != nil {
		return err
	}

	defer resp.Body.Close()

	LogReader(resp.Body)

	return D.TagImage(name+"%3A"+v, name)
}

// loadImage pulls a specified image into a docker instance
func (D *Docker) LoadImage(imagename string) (err error) {
	b := bytes.NewBuffer(nil)
	resp, err := D.h.Post("images/create?fromImage="+imagename,
		"text", b)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	LogReader(resp.Body)
	if resp.StatusCode != 200 {
		log.Printf("create status code is not 200 %s", resp.StatusCode)
	}

	log.Printf("Image fetched %s", imagename)
	return nil
}

// ListContainers gives the state for a specific docker container
func (D *Docker) ListContainers() (c []Container, err error) {
	resp, err := D.h.Get("containers/json")

	if err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return c, errors.New("Response code is not 200")
	}

	message, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	containers := &c
	err = json.Unmarshal(message, containers)
	return
}

// ListImages gives the state of the images for a specific Docker container
func (D *Docker) ListImages() (img []Image, err error) {
	resp, err := D.h.Get("images/json")

		if err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return img, errors.New("Response code is not 200")
	}

	message, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	images := &img
	err = json.Unmarshal(message, images)
	return
}
// ImageRunning states whether an image is running on a docker instance
func (D *Docker) ImageRunning(imagename string) (running bool, id string, err error) {

	containers, err := D.ListContainers()
	if err != nil {
		return false, "", err
	}

	for _, v := range containers {
		if strings.SplitN(v.Image, ":", 2)[0] == imagename {
			return true, v.Id, nil
		}
	}
	return
}
<|endoftext|>